feat(tools): add background embedding queue for thoughts
Some checks failed
CI / build-and-test (push) Failing after -29m22s

* Implement QueueThought method in BackfillTool for embedding generation
* Update CaptureTool to utilize embedding queuer for failed embeddings
* Add EmbeddingStatus field to Thought type for tracking embedding state
This commit is contained in:
2026-04-11 23:37:53 +02:00
parent 1d4dbad33f
commit 4d107cb87e
5 changed files with 61 additions and 13 deletions

View File

@@ -51,6 +51,30 @@ func NewBackfillTool(db *store.DB, provider ai.Provider, sessions *session.Activ
return &BackfillTool{store: db, provider: provider, sessions: sessions, logger: logger}
}
// QueueThought queues a single thought for background embedding generation.
// It is used by capture when the embedding provider is temporarily unavailable.
func (t *BackfillTool) QueueThought(ctx context.Context, id uuid.UUID, content string) {
go func() {
vec, err := t.provider.Embed(ctx, content)
if err != nil {
t.logger.Warn("background embedding retry failed",
slog.String("thought_id", id.String()),
slog.String("error", err.Error()),
)
return
}
model := t.provider.EmbeddingModel()
if err := t.store.UpsertEmbedding(ctx, id, model, vec); err != nil {
t.logger.Warn("background embedding upsert failed",
slog.String("thought_id", id.String()),
slog.String("error", err.Error()),
)
return
}
t.logger.Info("background embedding retry succeeded", slog.String("thought_id", id.String()))
}()
}
func (t *BackfillTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in BackfillInput) (*mcp.CallToolResult, BackfillOutput, error) {
limit := in.Limit
if limit <= 0 {

View File

@@ -6,6 +6,7 @@ import (
"strings"
"time"
"github.com/google/uuid"
"github.com/modelcontextprotocol/go-sdk/mcp"
"golang.org/x/sync/errgroup"
@@ -17,6 +18,11 @@ import (
thoughttypes "git.warky.dev/wdevs/amcs/internal/types"
)
// EmbeddingQueuer queues a thought for background embedding generation.
type EmbeddingQueuer interface {
QueueThought(ctx context.Context, id uuid.UUID, content string)
}
type CaptureTool struct {
store *store.DB
provider ai.Provider
@@ -24,6 +30,7 @@ type CaptureTool struct {
sessions *session.ActiveProjects
metadataTimeout time.Duration
retryer *MetadataRetryer
embedRetryer EmbeddingQueuer
log *slog.Logger
}
@@ -36,8 +43,8 @@ type CaptureOutput struct {
Thought thoughttypes.Thought `json:"thought"`
}
func NewCaptureTool(db *store.DB, provider ai.Provider, capture config.CaptureConfig, metadataTimeout time.Duration, sessions *session.ActiveProjects, retryer *MetadataRetryer, log *slog.Logger) *CaptureTool {
return &CaptureTool{store: db, provider: provider, capture: capture, sessions: sessions, metadataTimeout: metadataTimeout, retryer: retryer, log: log}
func NewCaptureTool(db *store.DB, provider ai.Provider, capture config.CaptureConfig, metadataTimeout time.Duration, sessions *session.ActiveProjects, retryer *MetadataRetryer, embedRetryer EmbeddingQueuer, log *slog.Logger) *CaptureTool {
return &CaptureTool{store: db, provider: provider, capture: capture, sessions: sessions, metadataTimeout: metadataTimeout, retryer: retryer, embedRetryer: embedRetryer, log: log}
}
func (t *CaptureTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in CaptureInput) (*mcp.CallToolResult, CaptureOutput, error) {
@@ -54,12 +61,18 @@ func (t *CaptureTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in C
var embedding []float32
rawMetadata := metadata.Fallback(t.capture)
metadataNeedsRetry := false
embeddingNeedsRetry := false
group, groupCtx := errgroup.WithContext(ctx)
group.Go(func() error {
vector, err := t.provider.Embed(groupCtx, content)
if err != nil {
return err
t.log.Warn("embedding failed, thought will be saved without embedding",
slog.String("provider", t.provider.Name()),
slog.String("error", err.Error()),
)
embeddingNeedsRetry = true
return nil
}
embedding = vector
return nil
@@ -106,6 +119,9 @@ func (t *CaptureTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in C
if metadataNeedsRetry && t.retryer != nil {
t.retryer.QueueThought(created.ID)
}
if embeddingNeedsRetry && t.embedRetryer != nil {
t.embedRetryer.QueueThought(ctx, created.ID, content)
}
return nil, CaptureOutput{Thought: created}, nil
}