From cfc78e049312adaf18d94b2056c94e65864b56ee Mon Sep 17 00:00:00 2001 From: Hein Puth Date: Tue, 14 Jul 2026 12:44:53 +0200 Subject: [PATCH] feat(tools): expose retrieval_mode in query-based tool responses Add a retrieval_mode field ("semantic" or "text") to the output of all five query-driven tools so callers can distinguish vector search from Postgres full-text fallback without inspecting server logs. Tools updated: search_thoughts, recall_context, get_project_context, summarize_thoughts, and related_thoughts (semantic neighbours only; omitempty so field is absent when include_semantic is false or no query is provided for the non-mandatory-query tools). The shared semanticSearch helper now returns the mode as a third return value. All callers updated; fixes two compile errors left by the previous worker (summarize.go and links.go were not capturing the new return). Co-Authored-By: Claude Sonnet 4.6 --- README.md | 2 +- internal/tools/context.go | 18 ++++++---- internal/tools/links.go | 9 +++-- internal/tools/recall.go | 12 ++++--- internal/tools/retrieval.go | 18 +++++++--- internal/tools/retrieval_test.go | 60 ++++++++++++++++++++++++++++++++ internal/tools/search.go | 7 ++-- internal/tools/summarize.go | 11 +++--- 8 files changed, 109 insertions(+), 28 deletions(-) create mode 100644 internal/tools/retrieval_test.go diff --git a/README.md b/README.md index bb2de48..580e1af 100644 --- a/README.md +++ b/README.md @@ -478,7 +478,7 @@ metadata_retry: include_archived: false ``` -**Search fallback**: when no embeddings exist for the active model in scope, `search_thoughts`, `recall_context`, `get_project_context`, `summarize_thoughts`, and `related_thoughts` automatically fall back to Postgres full-text search so results are never silently empty. +**Search fallback**: when no embeddings exist for the active model in scope, `search_thoughts`, `recall_context`, `get_project_context`, `summarize_thoughts`, and `related_thoughts` automatically fall back to Postgres full-text search so results are never silently empty. All five tools include a `retrieval_mode` field in their response (`"semantic"` or `"text"`) so callers can see which path was taken. ## Client Setup diff --git a/internal/tools/context.go b/internal/tools/context.go index 69cd717..2ad7c2e 100644 --- a/internal/tools/context.go +++ b/internal/tools/context.go @@ -36,9 +36,10 @@ type ContextItem struct { } type ProjectContextOutput struct { - Project thoughttypes.Project `json:"project"` - Context string `json:"context"` - Items []ContextItem `json:"items"` + Project thoughttypes.Project `json:"project"` + Context string `json:"context"` + Items []ContextItem `json:"items"` + RetrievalMode string `json:"retrieval_mode,omitempty"` } func NewContextTool(db *store.DB, embeddings *ai.EmbeddingRunner, search config.SearchConfig, sessions *session.ActiveProjects) *ContextTool { @@ -70,12 +71,14 @@ func (t *ContextTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in P }) } + var retrievalMode string query := strings.TrimSpace(in.Query) if query != "" { - semantic, err := semanticSearch(ctx, t.store, t.embeddings, t.search, query, limit, t.search.DefaultThreshold, &project.NumericID, nil) + semantic, mode, err := semanticSearch(ctx, t.store, t.embeddings, t.search, query, limit, t.search.DefaultThreshold, &project.NumericID, nil) if err != nil { return nil, ProjectContextOutput{}, err } + retrievalMode = mode for _, result := range semantic { key := fmt.Sprint(result.ID) if _, ok := seen[key]; ok { @@ -100,8 +103,9 @@ func (t *ContextTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in P _ = t.store.TouchProject(ctx, project.NumericID) return nil, ProjectContextOutput{ - Project: *project, - Context: contextBlock, - Items: items, + Project: *project, + Context: contextBlock, + Items: items, + RetrievalMode: retrievalMode, }, nil } diff --git a/internal/tools/links.go b/internal/tools/links.go index 2e78478..159d644 100644 --- a/internal/tools/links.go +++ b/internal/tools/links.go @@ -45,7 +45,8 @@ type RelatedThought struct { } type RelatedOutput struct { - Related []RelatedThought `json:"related"` + Related []RelatedThought `json:"related"` + RetrievalMode string `json:"retrieval_mode,omitempty"` } func NewLinksTool(db *store.DB, embeddings *ai.EmbeddingRunner, search config.SearchConfig) *LinksTool { @@ -117,11 +118,13 @@ func (t *LinksTool) Related(ctx context.Context, _ *mcp.CallToolRequest, in Rela includeSemantic = *in.IncludeSemantic } + var retrievalMode string if includeSemantic { - semantic, err := semanticSearch(ctx, t.store, t.embeddings, t.search, thought.Content, t.search.DefaultLimit, t.search.DefaultThreshold, thought.ProjectID, &thought.GUID) + semantic, mode, err := semanticSearch(ctx, t.store, t.embeddings, t.search, thought.Content, t.search.DefaultLimit, t.search.DefaultThreshold, thought.ProjectID, &thought.GUID) if err != nil { return nil, RelatedOutput{}, err } + retrievalMode = mode for _, item := range semantic { key := fmt.Sprint(item.ID) if _, ok := seen[key]; ok { @@ -138,5 +141,5 @@ func (t *LinksTool) Related(ctx context.Context, _ *mcp.CallToolRequest, in Rela } } - return nil, RelatedOutput{Related: related}, nil + return nil, RelatedOutput{Related: related, RetrievalMode: retrievalMode}, nil } diff --git a/internal/tools/recall.go b/internal/tools/recall.go index d392098..b715aa9 100644 --- a/internal/tools/recall.go +++ b/internal/tools/recall.go @@ -27,8 +27,9 @@ type RecallInput struct { } type RecallOutput struct { - Context string `json:"context"` - Items []ContextItem `json:"items"` + Context string `json:"context"` + Items []ContextItem `json:"items"` + RetrievalMode string `json:"retrieval_mode,omitempty"` } func NewRecallTool(db *store.DB, embeddings *ai.EmbeddingRunner, search config.SearchConfig, sessions *session.ActiveProjects) *RecallTool { @@ -53,7 +54,7 @@ func (t *RecallTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in Re projectID = &project.NumericID } - semantic, err := semanticSearch(ctx, t.store, t.embeddings, t.search, query, limit, t.search.DefaultThreshold, projectID, nil) + semantic, mode, err := semanticSearch(ctx, t.store, t.embeddings, t.search, query, limit, t.search.DefaultThreshold, projectID, nil) if err != nil { return nil, RecallOutput{}, err } @@ -100,7 +101,8 @@ func (t *RecallTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in Re } return nil, RecallOutput{ - Context: formatContextBlock(header, lines), - Items: items, + Context: formatContextBlock(header, lines), + Items: items, + RetrievalMode: mode, }, nil } diff --git a/internal/tools/retrieval.go b/internal/tools/retrieval.go index 0fcb32a..de8154c 100644 --- a/internal/tools/retrieval.go +++ b/internal/tools/retrieval.go @@ -11,10 +11,16 @@ import ( thoughttypes "git.warky.dev/wdevs/amcs/internal/types" ) +const ( + RetrievalModeSemantic = "semantic" + RetrievalModeText = "text" +) + // semanticSearch runs vector similarity search if embeddings exist for the // primary embedding model in the given scope, otherwise falls back to Postgres // full-text search. Search always uses the primary model so query vectors // match rows stored under the primary model name. +// It returns the results and the retrieval mode used ("semantic" or "text"). func semanticSearch( ctx context.Context, db *store.DB, @@ -25,20 +31,22 @@ func semanticSearch( threshold float64, projectID *int64, excludeID *uuid.UUID, -) ([]thoughttypes.SearchResult, error) { +) ([]thoughttypes.SearchResult, string, error) { model := embeddings.PrimaryModel() hasEmbeddings, err := db.HasEmbeddingsForModel(ctx, model, projectID) if err != nil { - return nil, err + return nil, "", err } if hasEmbeddings { embedding, err := embeddings.EmbedPrimary(ctx, query) if err != nil { - return nil, err + return nil, "", err } - return db.SearchSimilarThoughts(ctx, embedding, model, threshold, limit, projectID, excludeID) + results, err := db.SearchSimilarThoughts(ctx, embedding, model, threshold, limit, projectID, excludeID) + return results, RetrievalModeSemantic, err } - return db.SearchThoughtsText(ctx, query, limit, projectID, excludeID) + results, err := db.SearchThoughtsText(ctx, query, limit, projectID, excludeID) + return results, RetrievalModeText, err } diff --git a/internal/tools/retrieval_test.go b/internal/tools/retrieval_test.go new file mode 100644 index 0000000..68b5b80 --- /dev/null +++ b/internal/tools/retrieval_test.go @@ -0,0 +1,60 @@ +package tools + +import "testing" + +func TestRetrievalModeConstants(t *testing.T) { + if RetrievalModeSemantic != "semantic" { + t.Fatalf("RetrievalModeSemantic = %q, want %q", RetrievalModeSemantic, "semantic") + } + if RetrievalModeText != "text" { + t.Fatalf("RetrievalModeText = %q, want %q", RetrievalModeText, "text") + } +} + +func TestSearchOutputIncludesRetrievalMode(t *testing.T) { + out := SearchOutput{RetrievalMode: RetrievalModeSemantic} + if out.RetrievalMode != RetrievalModeSemantic { + t.Fatalf("SearchOutput.RetrievalMode = %q, want %q", out.RetrievalMode, RetrievalModeSemantic) + } +} + +func TestRelatedOutputIncludesRetrievalMode(t *testing.T) { + out := RelatedOutput{RetrievalMode: RetrievalModeText} + if out.RetrievalMode != RetrievalModeText { + t.Fatalf("RelatedOutput.RetrievalMode = %q, want %q", out.RetrievalMode, RetrievalModeText) + } + // retrieval_mode is omitempty — empty string means semantic search was skipped + out2 := RelatedOutput{} + if out2.RetrievalMode != "" { + t.Fatalf("RelatedOutput.RetrievalMode = %q, want empty when include_semantic is false", out2.RetrievalMode) + } +} + +func TestSummarizeOutputIncludesRetrievalModeOnlyWhenQueryUsed(t *testing.T) { + withQuery := SummarizeOutput{Summary: "s", Count: 1, RetrievalMode: RetrievalModeSemantic} + if withQuery.RetrievalMode != RetrievalModeSemantic { + t.Fatalf("SummarizeOutput.RetrievalMode = %q, want %q", withQuery.RetrievalMode, RetrievalModeSemantic) + } + withoutQuery := SummarizeOutput{Summary: "s", Count: 1} + if withoutQuery.RetrievalMode != "" { + t.Fatalf("SummarizeOutput.RetrievalMode = %q, want empty when no query", withoutQuery.RetrievalMode) + } +} + +func TestProjectContextOutputIncludesRetrievalModeOnlyWhenQueryUsed(t *testing.T) { + withQuery := ProjectContextOutput{RetrievalMode: RetrievalModeText} + if withQuery.RetrievalMode != RetrievalModeText { + t.Fatalf("ProjectContextOutput.RetrievalMode = %q, want %q", withQuery.RetrievalMode, RetrievalModeText) + } + withoutQuery := ProjectContextOutput{} + if withoutQuery.RetrievalMode != "" { + t.Fatalf("ProjectContextOutput.RetrievalMode = %q, want empty when no query", withoutQuery.RetrievalMode) + } +} + +func TestRecallOutputIncludesRetrievalMode(t *testing.T) { + out := RecallOutput{RetrievalMode: RetrievalModeSemantic} + if out.RetrievalMode != RetrievalModeSemantic { + t.Fatalf("RecallOutput.RetrievalMode = %q, want %q", out.RetrievalMode, RetrievalModeSemantic) + } +} diff --git a/internal/tools/search.go b/internal/tools/search.go index 7fa79be..8ce5e92 100644 --- a/internal/tools/search.go +++ b/internal/tools/search.go @@ -28,7 +28,8 @@ type SearchInput struct { } type SearchOutput struct { - Results []thoughttypes.SearchResult `json:"results"` + Results []thoughttypes.SearchResult `json:"results"` + RetrievalMode string `json:"retrieval_mode,omitempty"` } func NewSearchTool(db *store.DB, embeddings *ai.EmbeddingRunner, search config.SearchConfig, sessions *session.ActiveProjects) *SearchTool { @@ -55,10 +56,10 @@ func (t *SearchTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in Se _ = t.store.TouchProject(ctx, project.NumericID) } - results, err := semanticSearch(ctx, t.store, t.embeddings, t.search, query, limit, threshold, projectID, nil) + results, mode, err := semanticSearch(ctx, t.store, t.embeddings, t.search, query, limit, threshold, projectID, nil) if err != nil { return nil, SearchOutput{}, err } - return nil, SearchOutput{Results: results}, nil + return nil, SearchOutput{Results: results, RetrievalMode: mode}, nil } diff --git a/internal/tools/summarize.go b/internal/tools/summarize.go index 805c250..2887e0e 100644 --- a/internal/tools/summarize.go +++ b/internal/tools/summarize.go @@ -28,8 +28,9 @@ type SummarizeInput struct { } type SummarizeOutput struct { - Summary string `json:"summary"` - Count int `json:"count"` + Summary string `json:"summary"` + Count int `json:"count"` + RetrievalMode string `json:"retrieval_mode,omitempty"` } func NewSummarizeTool(db *store.DB, embeddings *ai.EmbeddingRunner, metadata *ai.MetadataRunner, search config.SearchConfig, sessions *session.ActiveProjects) *SummarizeTool { @@ -47,15 +48,17 @@ func (t *SummarizeTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in lines := make([]string, 0, limit) count := 0 + var retrievalMode string if query != "" { var projectID *int64 if project != nil { projectID = &project.NumericID } - results, err := semanticSearch(ctx, t.store, t.embeddings, t.search, query, limit, t.search.DefaultThreshold, projectID, nil) + results, mode, err := semanticSearch(ctx, t.store, t.embeddings, t.search, query, limit, t.search.DefaultThreshold, projectID, nil) if err != nil { return nil, SummarizeOutput{}, err } + retrievalMode = mode for i, result := range results { lines = append(lines, thoughtContextLine(i, result.Content, result.Metadata, result.Similarity)) } @@ -85,5 +88,5 @@ func (t *SummarizeTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in _ = t.store.TouchProject(ctx, project.NumericID) } - return nil, SummarizeOutput{Summary: summary, Count: count}, nil + return nil, SummarizeOutput{Summary: summary, Count: count, RetrievalMode: retrievalMode}, nil }