From 017052a8afe4900a52383a27117c88264d28d6bc Mon Sep 17 00:00:00 2001 From: xuzhaonan Date: Fri, 17 Jul 2026 16:16:49 +0800 Subject: [PATCH] feat(adk): replace forced tool call with response_format fallback for topic selection --- adk/middlewares/automemory/automemory.go | 208 +++++++++--- adk/middlewares/automemory/automemory_test.go | 301 +++++++++++++++++- adk/middlewares/automemory/consts.go | 27 +- adk/middlewares/automemory/prompt.go | 19 ++ adk/middlewares/automemory/utils.go | 144 +++++++-- adk/middlewares/automemory/utils_test.go | 127 ++++++++ components/model/option.go | 12 + schema/response_format.go | 43 +++ 8 files changed, 801 insertions(+), 80 deletions(-) create mode 100644 schema/response_format.go diff --git a/adk/middlewares/automemory/automemory.go b/adk/middlewares/automemory/automemory.go index 0371afb35..a0ce3faf7 100644 --- a/adk/middlewares/automemory/automemory.go +++ b/adk/middlewares/automemory/automemory.go @@ -73,10 +73,14 @@ type Config[M adk.MessageType] struct { // Optional. Defaults to a local in-process coordinator. Coordination *CoordinationConfig[M] - // OnError is called when automemory encounters an error. Errors are best-effort by default: - // the middleware will skip memory injection and allow the agent to continue. - // Optional. - OnError func(ctx context.Context, stage ErrorStage, err error) + // OnError is called when automemory encounters an error at a specific stage. + // Return a non-nil error to block the middleware and propagate the error to the caller. + // Return nil to skip the error and continue with best-effort degradation. + // Note: for stages that fire in background goroutines, the return value is ignored + // (see ErrorStage documentation for the full list). + // Optional. Defaults to defaultOnError, which blocks on render_instruction and + // snapshot_marshal, and logs all other errors without blocking. + OnError func(ctx context.Context, stage ErrorStage, err error) error } type ReadMode string @@ -150,8 +154,25 @@ type TopicSelectionConfig struct { // MaxTotalBytes caps the total rendered topic memory reminder. // Optional. Defaults to 16k. MaxTotalBytes int + + // OutputMode constrains the response format used for topic selection. + // Supported values: "json_schema", "json_object". + // When set, the fallback step uses only this response format (no further fallback). + // When empty (default), the middleware first tries a plain call (tools configured, + // no response_format), then falls back to json_schema and json_object in order. + OutputMode TopicSelectionOutputMode } +// TopicSelectionOutputMode specifies the response format used for topic selection fallback. +type TopicSelectionOutputMode string + +const ( + // TopicSelectionOutputModeJSONSchema uses response_format=json_schema for structured output. + TopicSelectionOutputModeJSONSchema TopicSelectionOutputMode = "json_schema" + // TopicSelectionOutputModeJSONObject uses response_format=json_object for structured output. + TopicSelectionOutputModeJSONObject TopicSelectionOutputMode = "json_object" +) + type WriteMode string const ( @@ -289,28 +310,30 @@ func (m *middleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAg // Sync distributed write cursor back into message extras so later runs on other // machines still carry a transcript-local marker. if nRunCtx.AgentInput != nil && len(nRunCtx.AgentInput.Messages) > 0 && m.coordination != nil && m.coordination.Coordinator != nil { - if sessionID, err := m.resolveSessionID(ctx, &adk.TypedChatModelAgentState[M]{Messages: nRunCtx.AgentInput.Messages}); err == nil && sessionID != "" { - localCursor := getWriteCursorFromMessages(nRunCtx.AgentInput.Messages) - coordKey := m.coordinatorKey(sessionID) - if remoteCursor, ok, err := getCoordinatorCursor(ctx, m.coordination.Coordinator, coordKey); err == nil && ok && remoteCursor > localCursor { - st := markWriteCursor(&adk.TypedChatModelAgentState[M]{Messages: nRunCtx.AgentInput.Messages}, remoteCursor) - if st != nil { - nRunCtx.AgentInput = &adk.TypedAgentInput[M]{ - Messages: st.Messages, - EnableStreaming: nRunCtx.AgentInput.EnableStreaming, - } + sessionID := m.coordination.SessionID + localCursor := getWriteCursorFromMessages(nRunCtx.AgentInput.Messages) + coordKey := m.coordinatorKey(sessionID) + if remoteCursor, ok, err := getCoordinatorCursor(ctx, m.coordination.Coordinator, coordKey); err == nil && ok && remoteCursor > localCursor { + st := markWriteCursor(&adk.TypedChatModelAgentState[M]{Messages: nRunCtx.AgentInput.Messages}, remoteCursor) + if st != nil { + nRunCtx.AgentInput = &adk.TypedAgentInput[M]{ + Messages: st.Messages, + EnableStreaming: nRunCtx.AgentInput.EnableStreaming, } } } } - // 1) System prompt: inject stable auto memory instruction and directory manifest (best-effort). + // 1) System prompt: inject stable auto memory instruction and directory manifest. + // If this fails, skip all subsequent memory injection to avoid partial context. instruction, err := m.renderInstruction(ctx, nRunCtx.Instruction) if err != nil { - m.onErr(ctx, OnErrorStageRenderInstruction, err) - } else { - nRunCtx.Instruction = instruction + if blockErr := m.onErr(ctx, OnErrorStageRenderInstruction, err); blockErr != nil { + return ctx, runCtx, blockErr + } + return ctx, &nRunCtx, nil } + nRunCtx.Instruction = instruction if nRunCtx.AgentInput == nil || len(nRunCtx.AgentInput.Messages) == 0 { return ctx, &nRunCtx, nil @@ -322,9 +345,10 @@ func (m *middleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAg if !hasMemoryIndexInjected(nRunCtx.AgentInput.Messages) { indexMsg, err := m.buildMemoryIndexMessage(ctx) if err != nil { - m.onErr(ctx, OnErrorStageRenderInstruction, err) + if blockErr := m.onErr(ctx, OnErrorStageReadMemoryIndex, err); blockErr != nil { + return ctx, runCtx, blockErr + } } else if !isNilMessage(indexMsg) { - m.sendTopicMemoryEvent(ctx, nRunCtx.AgentInput.Messages, indexMsg) reminders = append(reminders, indexMsg) } } @@ -334,14 +358,18 @@ func (m *middleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAg m.cfg.Read.Mode == ReadModeSync && m.topicSelectionEnabled() { memMsg, err := m.selectAndBuildTopicMemoryMessage(ctx, nRunCtx.AgentInput) if err != nil { - m.onErr(ctx, OnErrorStageTopicSelectionSync, err) + if blockErr := m.onErr(ctx, OnErrorStageTopicSelectionSync, err); blockErr != nil { + return ctx, runCtx, blockErr + } } else if !isNilMessage(memMsg) { - m.sendTopicMemoryEvent(ctx, nRunCtx.AgentInput.Messages, memMsg) reminders = append(reminders, memMsg) } } if len(reminders) > 0 { + for _, r := range reminders { + m.sendTopicMemoryEvent(ctx, nRunCtx.AgentInput.Messages, r) + } msgs := insertMessagesBeforeLastUserQuery(nRunCtx.AgentInput.Messages, reminders) nRunCtx.AgentInput = &adk.TypedAgentInput[M]{Messages: msgs, EnableStreaming: nRunCtx.AgentInput.EnableStreaming} } @@ -414,7 +442,9 @@ func (m *middleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk. err := fut.err fut.mu.Unlock() if err != nil { - m.onErr(ctx, OnErrorStageTopicSelectionAsync, err) + if blockErr := m.onErr(ctx, OnErrorStageTopicSelectionAsync, err); blockErr != nil { + return ctx, state, blockErr + } } var msgs []M @@ -621,14 +651,44 @@ func (m *middleware[M]) selectTopicCandidates( return nil, err } - toolInfo := topicSelectionToolInfo() + valid := make(map[string]struct{}, len(relToBundle)) + for k := range relToBundle { + valid[k] = struct{}{} + } + + mode := m.cfg.Read.TopicSelection.OutputMode + if mode != "" { + // Fixed mode: skip plain call, directly use the configured response_format. + return m.selectTopicWithResponseFormat(ctx, mode, userMsg, valid, topK) + } + + // Step 1: plain call (tools configured, no forced choice, no response_format). + selected, plainErr := m.selectTopicPlain(ctx, userMsg, valid, topK) + if plainErr == nil { + return selected, nil + } + + // Step 2: auto fallback with response_format — try json_schema, then json_object. + selected, err = m.selectTopicWithResponseFormat(ctx, TopicSelectionOutputModeJSONSchema, userMsg, valid, topK) + if err == nil { + return selected, nil + } + return m.selectTopicWithResponseFormat(ctx, TopicSelectionOutputModeJSONObject, userMsg, valid, topK) +} + +// selectTopicPlain calls the model with tools configured but no forced choice and no response_format. +func (m *middleware[M]) selectTopicPlain( + ctx context.Context, + userMsg string, + valid map[string]struct{}, + topK int, +) ([]string, error) { respStream, err := m.topicSelectionModel.Stream( ctx, []M{ - makeSystemMsg[M](getTopicSelectionSystemPrompt()), + makeSystemMsg[M](getTopicSelectionSystemPrompt() + "\n\n" + getTopicSelectionJSONOutputHint()), makeUserMsg[M](userMsg), }, - makeToolChoiceForced[M](toolInfo.Name), ) if err != nil { return nil, err @@ -639,11 +699,67 @@ func (m *middleware[M]) selectTopicCandidates( return nil, err } - valid := make(map[string]struct{}, len(relToBundle)) - for k := range relToBundle { - valid[k] = struct{}{} + return m.parseTopicSelectionResponse(resp, valid, topK) +} + +// selectTopicWithResponseFormat calls the model with tools and the given response_format. +func (m *middleware[M]) selectTopicWithResponseFormat( + ctx context.Context, + mode TopicSelectionOutputMode, + userMsg string, + valid map[string]struct{}, + topK int, +) ([]string, error) { + var rfOpt model.Option + switch mode { + case TopicSelectionOutputModeJSONSchema: + rfOpt = model.WithResponseFormat(&schema.ResponseFormat{ + Type: schema.ResponseFormatTypeJSONSchema, + JSONSchema: &schema.ResponseFormatJSONSchema{ + Schema: topicSelectionJSONSchema(), + }, + }) + case TopicSelectionOutputModeJSONObject: + rfOpt = model.WithResponseFormat(&schema.ResponseFormat{ + Type: schema.ResponseFormatTypeJSONObject, + }) + default: + return nil, fmt.Errorf("unsupported topic selection output mode: %q", mode) } - selected, err := parseTopicSelectionFromToolCall(resp, valid) + + respStream, err := m.topicSelectionModel.Stream( + ctx, + []M{ + makeSystemMsg[M](getTopicSelectionSystemPrompt() + "\n\n" + getTopicSelectionJSONOutputHint()), + makeUserMsg[M](userMsg), + }, + rfOpt, + ) + if err != nil { + return nil, err + } + + resp, err := concatMessageStream(respStream) + if err != nil { + return nil, err + } + + return m.parseTopicSelectionResponse(resp, valid, topK) +} + +// parseTopicSelectionResponse tries to extract selected memories from a model response, +// checking tool_call first, then falling back to content JSON parsing. +func (m *middleware[M]) parseTopicSelectionResponse(resp M, valid map[string]struct{}, topK int) ([]string, error) { + // Try tool call first. + if selected, err := parseTopicSelectionFromToolCall(resp, valid); err == nil { + if len(selected) > topK { + return selected[:topK], nil + } + return selected, nil + } + + // Fall back to content JSON parsing. + selected, err := parseTopicSelectionFromContent[M](resp, valid) if err != nil { return nil, err } @@ -730,11 +846,7 @@ func (m *middleware[M]) AfterAgent(ctx context.Context, state *adk.TypedChatMode return ctx, nil } - sessionID, err := m.resolveSessionID(ctx, state) - if err != nil { - m.onErr(ctx, OnErrorStageResolveSessionID, err) - return ctx, nil - } + sessionID := m.coordination.SessionID coordKey := m.coordinatorKey(sessionID) cursor := getWriteCursorFromMessages(state.Messages) @@ -775,7 +887,9 @@ func (m *middleware[M]) AfterAgent(ctx context.Context, state *adk.TypedChatMode case WriteModeSync: end := len(state.Messages) if err := m.runMemoryExtractionAgent(ctx, state.Messages, cursor, state.ToolInfos); err != nil { - m.onErr(ctx, OnErrorStageMemoryWriteSync, err) + if blockErr := m.onErr(ctx, OnErrorStageMemoryWriteSync, err); blockErr != nil { + return ctx, blockErr + } return ctx, nil } if coordKey != "" { @@ -787,7 +901,9 @@ func (m *middleware[M]) AfterAgent(ctx context.Context, state *adk.TypedChatMode case WriteModeAsync: if coordKey == "" { if err := m.runMemoryExtractionAgent(ctx, state.Messages, cursor, state.ToolInfos); err != nil { - m.onErr(ctx, OnErrorStageMemoryWriteSync, err) + if blockErr := m.onErr(ctx, OnErrorStageMemoryWriteSync, err); blockErr != nil { + return ctx, blockErr + } return ctx, nil } state = markWriteCursor(state, len(state.Messages)) @@ -795,17 +911,21 @@ func (m *middleware[M]) AfterAgent(ctx context.Context, state *adk.TypedChatMode } snap, err := buildPendingSnapshot(state.Messages, cursor, state.ToolInfos) if err != nil { - m.onErr(ctx, OnErrorStageSnapshotMarshal, err) + if blockErr := m.onErr(ctx, OnErrorStageSnapshotMarshal, err); blockErr != nil { + return ctx, blockErr + } return ctx, nil } unlock, ok, err := m.coordination.Coordinator.AcquireLock(ctx, coordKey, m.coordination.LockTTL) if err != nil { - m.onErr(ctx, OnErrorStageAcquireExtractionLock, err) + if blockErr := m.onErr(ctx, OnErrorStageAcquireExtractionLock, err); blockErr != nil { + return ctx, blockErr + } return ctx, nil } if !ok { if err := setCoordinatorPendingSnapshot(ctx, m.coordination.Coordinator, coordKey, snap, m.coordination.LockTTL); err != nil { - m.onErr(ctx, OnErrorStageStashPendingSnapshot, err) + m.onErr(ctx, OnErrorStageStashPendingSnapshot, err) // nolint: errcheck } return ctx, nil } @@ -823,7 +943,7 @@ func (m *middleware[M]) runExtractionDrain(ctx context.Context, coordKey string, return } if err := unlock(ctx); err != nil { - m.onErr(ctx, OnErrorStageReleaseExtractionLock, err) + m.onErr(ctx, OnErrorStageReleaseExtractionLock, err) //nolint:errcheck } }() @@ -831,16 +951,16 @@ func (m *middleware[M]) runExtractionDrain(ctx context.Context, coordKey string, for current != nil { msgs, cursor, toolInfos, err := decodePendingSnapshot[M](current) if err != nil { - m.onErr(ctx, OnErrorStageDecodePendingSnapshot, err) + m.onErr(ctx, OnErrorStageDecodePendingSnapshot, err) //nolint:errcheck } else if err := m.runMemoryExtractionAgent(ctx, msgs, cursor, toolInfos); err != nil { - m.onErr(ctx, OnErrorStageMemoryWriteAsync, err) + m.onErr(ctx, OnErrorStageMemoryWriteAsync, err) //nolint:errcheck } else { _ = setCoordinatorCursor(ctx, m.coordination.Coordinator, coordKey, len(msgs)) } next, loadErr := popCoordinatorPendingSnapshot(ctx, m.coordination.Coordinator, coordKey) if loadErr != nil { - m.onErr(ctx, OnErrorStageLoadPendingSnapshot, loadErr) + m.onErr(ctx, OnErrorStageLoadPendingSnapshot, loadErr) //nolint:errcheck return } current = next diff --git a/adk/middlewares/automemory/automemory_test.go b/adk/middlewares/automemory/automemory_test.go index 7a6a991a2..84931f198 100644 --- a/adk/middlewares/automemory/automemory_test.go +++ b/adk/middlewares/automemory/automemory_test.go @@ -232,8 +232,9 @@ func TestMiddleware_IndexInjection_CustomInstructionErrorReportsRenderStage(t *t GenInstruction: func(ctx context.Context) (string, error) { return "", fmt.Errorf("custom instruction failed") }, - OnError: func(ctx context.Context, stage ErrorStage, err error) { + OnError: func(ctx context.Context, stage ErrorStage, err error) error { stages = append(stages, stage) + return nil }, }) require.NoError(t, err) @@ -249,6 +250,98 @@ func TestMiddleware_IndexInjection_CustomInstructionErrorReportsRenderStage(t *t require.Equal(t, []ErrorStage{OnErrorStageRenderInstruction}, stages) } +func TestMiddleware_BeforeAgent_RenderInstructionError_BlocksByDefault(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + GenInstruction: func(ctx context.Context) (string, error) { + return "", fmt.Errorf("instruction render failed") + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + } + + _, _, err = mw.BeforeAgent(ctx, runCtx) + require.Error(t, err) + require.Contains(t, err.Error(), "instruction render failed") +} + +func TestMiddleware_BeforeAgent_RenderInstructionError_SkipsMemoryInjection(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [notes.md](notes.md) - notes\n", now) + b.put("/mem/notes.md", "some notes", now) + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: &contentModel{out: `{"selected_memories":["notes.md"]}`}, + GenInstruction: func(ctx context.Context) (string, error) { + return "", fmt.Errorf("broken") + }, + OnError: func(_ context.Context, _ ErrorStage, _ error) error { + return nil + }, + Read: &ReadConfig[*schema.Message]{Mode: ReadModeSync}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Equal(t, "base", out.Instruction) + // No memory messages should be injected when instruction rendering fails + require.Equal(t, 1, len(out.AgentInput.Messages), "should not inject any memory messages") +} + +func TestMiddleware_BeforeAgent_OnErrorCanBlock(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [notes.md](notes.md) - notes\n", now) + b.put("/mem/notes.md", "---\nname: Notes\ndescription: misc notes\ntype: project\n---\n\nsome notes", now) + + // Model that returns non-parseable content to trigger topic selection error + mdl := &contentModel{out: "I cannot select"} + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: mdl, + Read: &ReadConfig[*schema.Message]{Mode: ReadModeSync}, + OnError: func(_ context.Context, stage ErrorStage, err error) error { + if stage == OnErrorStageTopicSelectionSync { + return fmt.Errorf("blocked: %w", err) + } + return nil + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + } + + _, _, err = mw.BeforeAgent(ctx, runCtx) + require.Error(t, err) + require.Contains(t, err.Error(), "blocked:") +} + func TestNew_DoesNotMutateConfig(t *testing.T) { ctx := context.Background() b := NewInMemoryBackend() @@ -766,8 +859,9 @@ func TestMiddleware_AfterAgent_SyncExtractionWritesMemoryFiles(t *testing.T) { Mode: WriteModeSync, Model: extModel, }, - OnError: func(ctx context.Context, stage ErrorStage, err error) { + OnError: func(ctx context.Context, stage ErrorStage, err error) error { onErrStages = append(onErrStages, stage) + return nil }, }) require.NoError(t, err) @@ -876,8 +970,9 @@ func TestMiddleware_AfterAgent_SyncExtractionWritesMemoryDirectory(t *testing.T) Mode: WriteModeSync, Model: extModel, }, - OnError: func(ctx context.Context, stage ErrorStage, err error) { + OnError: func(ctx context.Context, stage ErrorStage, err error) error { onErrStages = append(onErrStages, stage) + return nil }, }) require.NoError(t, err) @@ -1578,3 +1673,203 @@ func TestMiddleware_AfterAgent_AsyncSetsPendingSnapshotWhenLockHeld(t *testing.T require.NoError(t, err) require.Equal(t, "remember pending", topic.Content) } + +// contentModel always returns the given content as a text response. +type contentModel struct { + out string +} + +func (m *contentModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return &schema.Message{Role: schema.Assistant, Content: m.out}, nil +} + +func (m *contentModel) Stream(_ context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(context.Background(), input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +// noResponseFormatModel rejects calls with ResponseFormat set. +type noResponseFormatModel struct{} + +func (m *noResponseFormatModel) Generate(_ context.Context, _ []*schema.Message, opts ...model.Option) (*schema.Message, error) { + common := model.GetCommonOptions(nil, opts...) + if common.ResponseFormat != nil { + return nil, fmt.Errorf("response_format not supported") + } + return &schema.Message{Role: schema.Assistant, Content: "no structured output"}, nil +} + +func (m *noResponseFormatModel) Stream(_ context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(context.Background(), input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func TestMiddleware_TopicSelection_FallbackToJSONContent(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md) - notes\n", now) + b.put("/mem/debugging.md", "---\nname: Debugging\ndescription: build and test commands\ntype: project\n---\n\n# Debugging\npnpm test\n", now) + + mdl := &contentModel{out: `{"selected_memories":["debugging.md"]}`} + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: mdl, + Read: &ReadConfig[*schema.Message]{Mode: ReadModeSync}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to run tests?")}}, + } + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + topicCount := countTopicMemoryMessages(out.AgentInput.Messages) + require.Equal(t, 1, topicCount, "should inject topic memory via fallback") +} + +func TestMiddleware_TopicSelection_FixedOutputMode(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md) - notes\n", now) + b.put("/mem/debugging.md", "---\nname: Debugging\ndescription: build and test commands\ntype: project\n---\n\n# Debugging\npnpm test\n", now) + + mdl := &contentModel{out: `{"selected_memories":["debugging.md"]}`} + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: mdl, + Read: &ReadConfig[*schema.Message]{ + Mode: ReadModeSync, + TopicSelection: &TopicSelectionConfig{ + OutputMode: TopicSelectionOutputModeJSONObject, + }, + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to run tests?")}}, + } + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + topicCount := countTopicMemoryMessages(out.AgentInput.Messages) + require.Equal(t, 1, topicCount, "should inject topic memory with fixed json_object mode") +} + +func TestMiddleware_TopicSelection_FixedOutputModeFails(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md) - notes\n", now) + b.put("/mem/debugging.md", "---\nname: Debugging\ndescription: build and test commands\ntype: project\n---\n\n# Debugging\npnpm test\n", now) + + mdl := &noResponseFormatModel{} + + var capturedErr error + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: mdl, + Read: &ReadConfig[*schema.Message]{ + Mode: ReadModeSync, + TopicSelection: &TopicSelectionConfig{ + OutputMode: TopicSelectionOutputModeJSONSchema, + }, + }, + OnError: func(_ context.Context, _ ErrorStage, err error) error { + capturedErr = err + return nil + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to run tests?")}}, + } + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + topicCount := countTopicMemoryMessages(out.AgentInput.Messages) + require.Equal(t, 0, topicCount, "should not inject topic memory when fixed output mode fails") + require.NotNil(t, capturedErr, "OnError should be called when output mode fails") +} + + +func TestMiddleware_TopicSelection_FallbackJSONSchemaToJSONObject(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md) - notes\n", now) + b.put("/mem/debugging.md", "---\nname: Debugging\ndescription: build and test commands\ntype: project\n---\n\n# Debugging\npnpm test\n", now) + + // Model returns non-JSON for plain call, rejects json_schema, returns JSON for json_object. + // Step 1 (plain) fails → step 2 json_schema fails → step 2 json_object succeeds. + mdl := &selectiveResponseModel{ + plainOut: "Let me think...", + jsonObjectOut: `{"selected_memories":["debugging.md"]}`, + } + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: mdl, + Read: &ReadConfig[*schema.Message]{Mode: ReadModeSync}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to run tests?")}}, + } + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + topicCount := countTopicMemoryMessages(out.AgentInput.Messages) + require.Equal(t, 1, topicCount, "should inject topic memory via json_object fallback") +} + +// selectiveResponseModel rejects json_schema, returns non-parseable for plain, returns JSON for json_object. +type selectiveResponseModel struct { + plainOut string + jsonObjectOut string +} + +func (m *selectiveResponseModel) Generate(_ context.Context, _ []*schema.Message, opts ...model.Option) (*schema.Message, error) { + common := model.GetCommonOptions(nil, opts...) + if common.ResponseFormat != nil { + switch common.ResponseFormat.Type { + case schema.ResponseFormatTypeJSONSchema: + return nil, fmt.Errorf("json_schema not supported") + case schema.ResponseFormatTypeJSONObject: + return &schema.Message{Role: schema.Assistant, Content: m.jsonObjectOut}, nil + } + } + return &schema.Message{Role: schema.Assistant, Content: m.plainOut}, nil +} + +func (m *selectiveResponseModel) Stream(_ context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(context.Background(), input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} diff --git a/adk/middlewares/automemory/consts.go b/adk/middlewares/automemory/consts.go index f5ee3aa43..3675bfb82 100644 --- a/adk/middlewares/automemory/consts.go +++ b/adk/middlewares/automemory/consts.go @@ -41,13 +41,34 @@ const ( // ErrorStage error stage during auto memory processing type ErrorStage string -// OnError stage constants. These values are stable identifiers used to report -// best-effort failures through Config.OnError. +// OnError stage constants identify where an error occurred in the automemory lifecycle. +// +// The OnError callback's return value controls whether the middleware blocks (non-nil) +// or continues with degraded behavior (nil). However, for stages that fire in background +// goroutines or non-propagatable contexts, the return value is ignored — the error is +// reported for observability only. +// +// Stages where return value is honored (can block): +// - OnErrorStageRenderInstruction (BeforeAgent) +// - OnErrorStageReadMemoryIndex (BeforeAgent) +// - OnErrorStageTopicSelectionSync (BeforeAgent) +// - OnErrorStageTopicSelectionAsync (BeforeModelRewriteState) +// - OnErrorStageMemoryWriteSync (AfterAgent) +// - OnErrorStageSnapshotMarshal (AfterAgent) +// - OnErrorStageAcquireExtractionLock (AfterAgent) +// +// Stages where return value is ignored (background/async, notification only): +// - OnErrorStageStashPendingSnapshot +// - OnErrorStageReleaseExtractionLock +// - OnErrorStageDecodePendingSnapshot +// - OnErrorStageMemoryWriteAsync +// - OnErrorStageLoadPendingSnapshot +// - OnErrorStageSendSessionEvent const ( OnErrorStageTopicSelectionSync ErrorStage = "topic_selection_sync" OnErrorStageTopicSelectionAsync ErrorStage = "topic_selection_async" OnErrorStageRenderInstruction ErrorStage = "render_instruction" - OnErrorStageResolveSessionID ErrorStage = "resolve_session_id" + OnErrorStageReadMemoryIndex ErrorStage = "read_memory_index" OnErrorStageMemoryWriteSync ErrorStage = "memory_write_sync" OnErrorStageSnapshotMarshal ErrorStage = "snapshot_marshal" OnErrorStageAcquireExtractionLock ErrorStage = "acquire_extraction_lock" diff --git a/adk/middlewares/automemory/prompt.go b/adk/middlewares/automemory/prompt.go index 434fd9d95..9a2e09da4 100644 --- a/adk/middlewares/automemory/prompt.go +++ b/adk/middlewares/automemory/prompt.go @@ -146,6 +146,18 @@ Recently used tools: defaultTopicMemoryTruncNotifyChinese = ` > 该记忆文件已被截断（{reason}）。请使用 Read 工具查看完整文件：{abs_path}` + + defaultTopicSelectionJSONOutputHint = `You must respond with a JSON object in the following format (no additional text, no markdown code block): +{"selected_memories": ["path/to/memory1.md", "path/to/memory2.md"]} + +Example — if the available memories include "debugging.md" and "patterns.md", and only "debugging.md" is relevant: +{"selected_memories": ["debugging.md"]}` + + defaultTopicSelectionJSONOutputHintChinese = `你必须以如下 JSON 格式回复（不要附加任何额外文本，不要使用 markdown 代码块）： +{"selected_memories": ["path/to/memory1.md", "path/to/memory2.md"]} + +示例 — 假设可用记忆包含 "debugging.md" 和 "patterns.md"，且仅 "debugging.md" 相关： +{"selected_memories": ["debugging.md"]}` ) type memoryIndexPromptInfo struct { @@ -456,6 +468,13 @@ func getTopicSelectionUserPrompt() string { }) } +func getTopicSelectionJSONOutputHint() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: defaultTopicSelectionJSONOutputHint, + Chinese: defaultTopicSelectionJSONOutputHintChinese, + }) +} + func getTopicMemoryTruncNotify() string { return internal.SelectPrompt(internal.I18nPrompts{ English: defaultTopicMemoryTruncNotify, diff --git a/adk/middlewares/automemory/utils.go b/adk/middlewares/automemory/utils.go index b198ee52d..2ae44a3ca 100644 --- a/adk/middlewares/automemory/utils.go +++ b/adk/middlewares/automemory/utils.go @@ -21,16 +21,17 @@ import ( "encoding/json" "fmt" "io" + "log" "path/filepath" "sort" "strings" "time" + "github.com/eino-contrib/jsonschema" "gopkg.in/yaml.v3" "github.com/cloudwego/eino/adk" adkfs "github.com/cloudwego/eino/adk/middlewares/filesystem" - "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" ) @@ -247,6 +248,16 @@ func topicSelectionToolInfo() *schema.ToolInfo { } } +func topicSelectionJSONSchema() *jsonschema.Schema { + r := &jsonschema.Reflector{ + Anonymous: true, + DoNotReference: true, + } + s := r.Reflect(&topicSelectionResp{}) + s.Version = "" + return s +} + func parseTopicSelectionFromToolCall[M adk.MessageType](msg M, valid map[string]struct{}) ([]string, error) { toolCalls := messageToolCalls(msg) if len(toolCalls) == 0 { @@ -388,6 +399,87 @@ func userMessageTextContent[M adk.MessageType](msg M) string { } } +func assistantTextContent[M adk.MessageType](msg M) string { + switch m := any(msg).(type) { + case *schema.Message: + if m == nil { + return "" + } + return m.Content + case *schema.AgenticMessage: + if m == nil { + return "" + } + parts := make([]string, 0, len(m.ContentBlocks)) + for _, block := range m.ContentBlocks { + if block != nil && block.AssistantGenText != nil { + parts = append(parts, block.AssistantGenText.Text) + } + } + return strings.Join(parts, "\n") + default: + panic("unreachable") + } +} + +func parseTopicSelectionFromContent[M adk.MessageType](msg M, valid map[string]struct{}) ([]string, error) { + content := assistantTextContent[M](msg) + if content == "" { + return nil, fmt.Errorf("empty response content") + } + + jsonStr := extractJSON(content) + if jsonStr == "" { + return nil, fmt.Errorf("no JSON found in response content") + } + + var parsed topicSelectionResp + if err := json.Unmarshal([]byte(jsonStr), &parsed); err != nil { + return nil, fmt.Errorf("failed to parse JSON from content: %w", err) + } + + out := normalizeSelected(parsed.SelectedMemories) + filtered := make([]string, 0, len(out)) + for _, p := range out { + if _, ok := valid[p]; ok { + filtered = append(filtered, p) + } + } + return filtered, nil +} + +func extractJSON(s string) string { + s = strings.TrimSpace(s) + + if start := strings.Index(s, "```json"); start != -1 { + body := s[start+7:] + if end := strings.Index(body, "```"); end != -1 { + return strings.TrimSpace(body[:end]) + } + } + if start := strings.Index(s, "```"); start != -1 { + body := s[start+3:] + if nl := strings.IndexByte(body, '\n'); nl != -1 { + body = body[nl+1:] + } + if end := strings.Index(body, "```"); end != -1 { + candidate := strings.TrimSpace(body[:end]) + if len(candidate) > 0 && candidate[0] == '{' { + return candidate + } + } + } + + if idx := strings.IndexByte(s, '{'); idx != -1 { + candidate := s[idx:] + if last := strings.LastIndexByte(candidate, '}'); last != -1 { + return candidate[:last+1] + } + } + + return "" +} + func getMsgExtra[M adk.MessageType](msg M) map[string]any { switch m := any(msg).(type) { case *schema.Message: @@ -447,23 +539,6 @@ func makeSystemMsg[M adk.MessageType](text string) M { } } -func makeToolChoiceForced[M adk.MessageType](name string) model.Option { - var zero M - switch any(zero).(type) { - case *schema.Message: - return model.WithToolChoice(schema.ToolChoiceForced, name) - case *schema.AgenticMessage: - return model.WithAgenticToolChoice(&schema.AgenticToolChoice{ - Type: schema.ToolChoiceForced, - Forced: &schema.AgenticForcedToolChoice{ - Tools: []*schema.AllowedTool{{FunctionName: name}}, - }, - }) - default: - panic("unreachable") - } -} - func messageToolCalls[M adk.MessageType](msg M) []schema.ToolCall { switch m := any(msg).(type) { case *schema.Message: @@ -851,13 +926,29 @@ func topicSelectionConfigEnabled(cfg *TopicSelectionConfig) bool { return cfg != nil && cfg.Enable != nil && *cfg.Enable } -func (m *middleware[M]) onErr(ctx context.Context, stage ErrorStage, err error) { +// defaultOnError is the default error handler used when Config.OnError is nil. +// It blocks on instruction rendering failures and snapshot marshal errors (which indicate bugs), +// logs all other errors via log.Printf, and returns nil to allow degraded execution. +func defaultOnError(_ context.Context, stage ErrorStage, err error) error { + switch stage { + case OnErrorStageRenderInstruction, OnErrorStageSnapshotMarshal: + log.Printf("[automemory] fatal error at stage %q: %v", stage, err) + return err + default: + log.Printf("[automemory] non-fatal error at stage %q: %v", stage, err) + return nil + } +} + +func (m *middleware[M]) onErr(ctx context.Context, stage ErrorStage, err error) error { if err == nil { - return + return nil } - if m.cfg != nil && m.cfg.OnError != nil { - m.cfg.OnError(ctx, stage, err) + handler := m.cfg.OnError + if handler == nil { + handler = defaultOnError } + return handler(ctx, stage, err) } func (m *middleware[M]) lastUserMessage(agentIn *adk.TypedAgentInput[M]) (M, bool) { @@ -885,13 +976,6 @@ func (m *middleware[M]) topicSelectionTopK() int { return topK } -func (m *middleware[M]) resolveSessionID(ctx context.Context, state *adk.TypedChatModelAgentState[M]) (string, error) { - if m.coordination != nil { - return strings.TrimSpace(m.coordination.SessionID), nil - } - return "", nil -} - func (m *middleware[M]) sendTopicMemoryEvent(ctx context.Context, msgs []M, memMsg M) { var beforeID string if len(msgs) > 0 && !isNilMessage(msgs[len(msgs)-1]) { @@ -908,6 +992,6 @@ func (m *middleware[M]) sendTopicMemoryEvent(ctx context.Context, msgs []M, memM }, }, }); sendEventErr != nil { - m.onErr(ctx, OnErrorStageSendSessionEvent, sendEventErr) + _ = m.onErr(ctx, OnErrorStageSendSessionEvent, sendEventErr) //nolint:errcheck } } diff --git a/adk/middlewares/automemory/utils_test.go b/adk/middlewares/automemory/utils_test.go index 5ff7ff851..b56852496 100644 --- a/adk/middlewares/automemory/utils_test.go +++ b/adk/middlewares/automemory/utils_test.go @@ -88,3 +88,130 @@ func TestConcatMessageStream_WithToolCalls(t *testing.T) { require.Len(t, msg.ToolCalls, 1) assert.Equal(t, "search", msg.ToolCalls[0].Function.Name) } + +func TestParseTopicSelectionFromContent_PureJSON(t *testing.T) { + valid := map[string]struct{}{ + "debugging.md": {}, + "patterns.md": {}, + "notes.md": {}, + } + msg := &schema.Message{ + Role: schema.Assistant, + Content: `{"selected_memories": ["debugging.md", "patterns.md"]}`, + } + + selected, err := parseTopicSelectionFromContent[*schema.Message](msg, valid) + require.NoError(t, err) + assert.Equal(t, []string{"debugging.md", "patterns.md"}, selected) +} + +func TestParseTopicSelectionFromContent_MarkdownCodeBlock(t *testing.T) { + valid := map[string]struct{}{ + "debugging.md": {}, + "patterns.md": {}, + } + msg := &schema.Message{ + Role: schema.Assistant, + Content: "Here is my selection:\n```json\n{\"selected_memories\": [\"debugging.md\"]}\n```\n", + } + + selected, err := parseTopicSelectionFromContent[*schema.Message](msg, valid) + require.NoError(t, err) + assert.Equal(t, []string{"debugging.md"}, selected) +} + +func TestParseTopicSelectionFromContent_FiltersInvalidPaths(t *testing.T) { + valid := map[string]struct{}{ + "debugging.md": {}, + } + msg := &schema.Message{ + Role: schema.Assistant, + Content: `{"selected_memories": ["debugging.md", "nonexistent.md"]}`, + } + + selected, err := parseTopicSelectionFromContent[*schema.Message](msg, valid) + require.NoError(t, err) + assert.Equal(t, []string{"debugging.md"}, selected) +} + +func TestParseTopicSelectionFromContent_EmptyContent(t *testing.T) { + valid := map[string]struct{}{"a.md": {}} + msg := &schema.Message{Role: schema.Assistant, Content: ""} + + _, err := parseTopicSelectionFromContent[*schema.Message](msg, valid) + require.Error(t, err) + assert.Contains(t, err.Error(), "empty response content") +} + +func TestParseTopicSelectionFromContent_NoJSON(t *testing.T) { + valid := map[string]struct{}{"a.md": {}} + msg := &schema.Message{Role: schema.Assistant, Content: "I don't know what to pick."} + + _, err := parseTopicSelectionFromContent[*schema.Message](msg, valid) + require.Error(t, err) + assert.Contains(t, err.Error(), "no JSON found") +} + +func TestParseTopicSelectionFromContent_EmptySelection(t *testing.T) { + valid := map[string]struct{}{"a.md": {}} + msg := &schema.Message{ + Role: schema.Assistant, + Content: `{"selected_memories": []}`, + } + + selected, err := parseTopicSelectionFromContent[*schema.Message](msg, valid) + require.NoError(t, err) + assert.Empty(t, selected) +} + +func TestExtractJSON(t *testing.T) { + tests := []struct { + name string + input string + want string + }{ + { + name: "pure json", + input: `{"selected_memories": ["a.md"]}`, + want: `{"selected_memories": ["a.md"]}`, + }, + { + name: "json with prefix text", + input: "Here are the results:\n{\"selected_memories\": [\"b.md\"]}", + want: `{"selected_memories": ["b.md"]}`, + }, + { + name: "markdown json block", + input: "```json\n{\"selected_memories\": [\"c.md\"]}\n```", + want: `{"selected_memories": ["c.md"]}`, + }, + { + name: "markdown generic block with json", + input: "```\n{\"selected_memories\": [\"d.md\"]}\n```", + want: `{"selected_memories": ["d.md"]}`, + }, + { + name: "no json", + input: "I cannot select any memories.", + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := extractJSON(tt.input) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTopicSelectionJSONSchema(t *testing.T) { + s := topicSelectionJSONSchema() + require.NotNil(t, s) + assert.Equal(t, "object", s.Type) + props := s.Properties + require.NotNil(t, props) + val, ok := props.Get("selected_memories") + require.True(t, ok) + assert.Equal(t, "array", val.Type) +} diff --git a/components/model/option.go b/components/model/option.go index 2222e14a1..82f78b0d5 100644 --- a/components/model/option.go +++ b/components/model/option.go @@ -42,6 +42,8 @@ type Options struct { MaxTokens *int // Stop is the stop words for the model, which controls the stopping condition of the model. Stop []string + // ResponseFormat controls the structured output format of the model response. + ResponseFormat *schema.ResponseFormat // Options only available for chat model. @@ -173,6 +175,16 @@ func WithAgenticToolChoice(toolChoice *schema.AgenticToolChoice) Option { } } +// WithResponseFormat sets the response format for the model. +// Only available for ChatModel. +func WithResponseFormat(rf *schema.ResponseFormat) Option { + return Option{ + apply: func(opts *Options) { + opts.ResponseFormat = rf + }, + } +} + // WrapImplSpecificOptFn is the option to wrap the implementation specific option function. // WrapImplSpecificOptFn wraps an implementation-specific option function into // an [Option] so it can be passed alongside standard options. diff --git a/schema/response_format.go b/schema/response_format.go new file mode 100644 index 000000000..bdbc92f3d --- /dev/null +++ b/schema/response_format.go @@ -0,0 +1,43 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import "github.com/eino-contrib/jsonschema" + +// ResponseFormatType specifies the format the model must output. +type ResponseFormatType string + +const ( + // ResponseFormatTypeJSONObject forces the model to output valid JSON. + ResponseFormatTypeJSONObject ResponseFormatType = "json_object" + // ResponseFormatTypeJSONSchema forces the model to output JSON conforming to a given schema. + ResponseFormatTypeJSONSchema ResponseFormatType = "json_schema" +) + +// ResponseFormat controls the structured output format of the model response. +type ResponseFormat struct { + // Type specifies the response format type. + Type ResponseFormatType + // JSONSchema specifies the schema when Type is ResponseFormatJSONSchema. + JSONSchema *ResponseFormatJSONSchema +} + +// ResponseFormatJSONSchema defines the JSON Schema for structured output. +type ResponseFormatJSONSchema struct { + // Schema is the JSON Schema definition. + Schema *jsonschema.Schema +}