Expose retrieval mode in query responses #37
@@ -478,7 +478,7 @@ metadata_retry:
|
|||||||
include_archived: false
|
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
|
## Client Setup
|
||||||
|
|
||||||
|
|||||||
@@ -36,9 +36,10 @@ type ContextItem struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type ProjectContextOutput struct {
|
type ProjectContextOutput struct {
|
||||||
Project thoughttypes.Project `json:"project"`
|
Project thoughttypes.Project `json:"project"`
|
||||||
Context string `json:"context"`
|
Context string `json:"context"`
|
||||||
Items []ContextItem `json:"items"`
|
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 {
|
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)
|
query := strings.TrimSpace(in.Query)
|
||||||
if 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 {
|
if err != nil {
|
||||||
return nil, ProjectContextOutput{}, err
|
return nil, ProjectContextOutput{}, err
|
||||||
}
|
}
|
||||||
|
retrievalMode = mode
|
||||||
for _, result := range semantic {
|
for _, result := range semantic {
|
||||||
key := fmt.Sprint(result.ID)
|
key := fmt.Sprint(result.ID)
|
||||||
if _, ok := seen[key]; ok {
|
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)
|
_ = t.store.TouchProject(ctx, project.NumericID)
|
||||||
|
|
||||||
return nil, ProjectContextOutput{
|
return nil, ProjectContextOutput{
|
||||||
Project: *project,
|
Project: *project,
|
||||||
Context: contextBlock,
|
Context: contextBlock,
|
||||||
Items: items,
|
Items: items,
|
||||||
|
RetrievalMode: retrievalMode,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -45,7 +45,8 @@ type RelatedThought struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type RelatedOutput 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 {
|
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
|
includeSemantic = *in.IncludeSemantic
|
||||||
}
|
}
|
||||||
|
|
||||||
|
var retrievalMode string
|
||||||
if includeSemantic {
|
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 {
|
if err != nil {
|
||||||
return nil, RelatedOutput{}, err
|
return nil, RelatedOutput{}, err
|
||||||
}
|
}
|
||||||
|
retrievalMode = mode
|
||||||
for _, item := range semantic {
|
for _, item := range semantic {
|
||||||
key := fmt.Sprint(item.ID)
|
key := fmt.Sprint(item.ID)
|
||||||
if _, ok := seen[key]; ok {
|
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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -27,8 +27,9 @@ type RecallInput struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type RecallOutput struct {
|
type RecallOutput struct {
|
||||||
Context string `json:"context"`
|
Context string `json:"context"`
|
||||||
Items []ContextItem `json:"items"`
|
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 {
|
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
|
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 {
|
if err != nil {
|
||||||
return nil, RecallOutput{}, err
|
return nil, RecallOutput{}, err
|
||||||
}
|
}
|
||||||
@@ -100,7 +101,8 @@ func (t *RecallTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in Re
|
|||||||
}
|
}
|
||||||
|
|
||||||
return nil, RecallOutput{
|
return nil, RecallOutput{
|
||||||
Context: formatContextBlock(header, lines),
|
Context: formatContextBlock(header, lines),
|
||||||
Items: items,
|
Items: items,
|
||||||
|
RetrievalMode: mode,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,10 +11,16 @@ import (
|
|||||||
thoughttypes "git.warky.dev/wdevs/amcs/internal/types"
|
thoughttypes "git.warky.dev/wdevs/amcs/internal/types"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
RetrievalModeSemantic = "semantic"
|
||||||
|
RetrievalModeText = "text"
|
||||||
|
)
|
||||||
|
|
||||||
// semanticSearch runs vector similarity search if embeddings exist for the
|
// semanticSearch runs vector similarity search if embeddings exist for the
|
||||||
// primary embedding model in the given scope, otherwise falls back to Postgres
|
// primary embedding model in the given scope, otherwise falls back to Postgres
|
||||||
// full-text search. Search always uses the primary model so query vectors
|
// full-text search. Search always uses the primary model so query vectors
|
||||||
// match rows stored under the primary model name.
|
// match rows stored under the primary model name.
|
||||||
|
// It returns the results and the retrieval mode used ("semantic" or "text").
|
||||||
func semanticSearch(
|
func semanticSearch(
|
||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
db *store.DB,
|
db *store.DB,
|
||||||
@@ -25,20 +31,22 @@ func semanticSearch(
|
|||||||
threshold float64,
|
threshold float64,
|
||||||
projectID *int64,
|
projectID *int64,
|
||||||
excludeID *uuid.UUID,
|
excludeID *uuid.UUID,
|
||||||
) ([]thoughttypes.SearchResult, error) {
|
) ([]thoughttypes.SearchResult, string, error) {
|
||||||
model := embeddings.PrimaryModel()
|
model := embeddings.PrimaryModel()
|
||||||
hasEmbeddings, err := db.HasEmbeddingsForModel(ctx, model, projectID)
|
hasEmbeddings, err := db.HasEmbeddingsForModel(ctx, model, projectID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, "", err
|
||||||
}
|
}
|
||||||
|
|
||||||
if hasEmbeddings {
|
if hasEmbeddings {
|
||||||
embedding, err := embeddings.EmbedPrimary(ctx, query)
|
embedding, err := embeddings.EmbedPrimary(ctx, query)
|
||||||
if err != nil {
|
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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -28,7 +28,8 @@ type SearchInput struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type SearchOutput 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 {
|
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)
|
_ = 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 {
|
if err != nil {
|
||||||
return nil, SearchOutput{}, err
|
return nil, SearchOutput{}, err
|
||||||
}
|
}
|
||||||
|
|
||||||
return nil, SearchOutput{Results: results}, nil
|
return nil, SearchOutput{Results: results, RetrievalMode: mode}, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -28,8 +28,9 @@ type SummarizeInput struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type SummarizeOutput struct {
|
type SummarizeOutput struct {
|
||||||
Summary string `json:"summary"`
|
Summary string `json:"summary"`
|
||||||
Count int `json:"count"`
|
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 {
|
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)
|
lines := make([]string, 0, limit)
|
||||||
count := 0
|
count := 0
|
||||||
|
|
||||||
|
var retrievalMode string
|
||||||
if query != "" {
|
if query != "" {
|
||||||
var projectID *int64
|
var projectID *int64
|
||||||
if project != nil {
|
if project != nil {
|
||||||
projectID = &project.NumericID
|
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 {
|
if err != nil {
|
||||||
return nil, SummarizeOutput{}, err
|
return nil, SummarizeOutput{}, err
|
||||||
}
|
}
|
||||||
|
retrievalMode = mode
|
||||||
for i, result := range results {
|
for i, result := range results {
|
||||||
lines = append(lines, thoughtContextLine(i, result.Content, result.Metadata, result.Similarity))
|
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)
|
_ = t.store.TouchProject(ctx, project.NumericID)
|
||||||
}
|
}
|
||||||
|
|
||||||
return nil, SummarizeOutput{Summary: summary, Count: count}, nil
|
return nil, SummarizeOutput{Summary: summary, Count: count, RetrievalMode: retrievalMode}, nil
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user