feat(tools): add background embedding queue for thoughts
Some checks failed
CI / build-and-test (push) Failing after -29m22s
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:
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user