130 lines
4.1 KiB
Go
130 lines
4.1 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
|
|
"git.warky.dev/wdevs/amcs/internal/ai"
|
|
"git.warky.dev/wdevs/amcs/internal/config"
|
|
"git.warky.dev/wdevs/amcs/internal/metadata"
|
|
"git.warky.dev/wdevs/amcs/internal/session"
|
|
"git.warky.dev/wdevs/amcs/internal/store"
|
|
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
|
|
capture config.CaptureConfig
|
|
sessions *session.ActiveProjects
|
|
metadataTimeout time.Duration
|
|
retryer *MetadataRetryer
|
|
embedRetryer EmbeddingQueuer
|
|
log *slog.Logger
|
|
}
|
|
|
|
type CaptureInput struct {
|
|
Content string `json:"content" jsonschema:"the thought or note to capture"`
|
|
Project string `json:"project,omitempty" jsonschema:"optional project name or id to associate with the thought"`
|
|
}
|
|
|
|
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, 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) {
|
|
content := strings.TrimSpace(in.Content)
|
|
if content == "" {
|
|
return nil, CaptureOutput{}, errRequiredField("content")
|
|
}
|
|
|
|
project, err := resolveProject(ctx, t.store, t.sessions, req, in.Project, false)
|
|
if err != nil {
|
|
return nil, CaptureOutput{}, err
|
|
}
|
|
|
|
rawMetadata := metadata.Fallback(t.capture)
|
|
thought := thoughttypes.Thought{
|
|
Content: content,
|
|
Metadata: rawMetadata,
|
|
}
|
|
if project != nil {
|
|
thought.ProjectID = &project.ID
|
|
}
|
|
|
|
created, err := t.store.InsertThought(ctx, thought, t.provider.EmbeddingModel())
|
|
if err != nil {
|
|
return nil, CaptureOutput{}, err
|
|
}
|
|
if project != nil {
|
|
_ = t.store.TouchProject(ctx, project.ID)
|
|
}
|
|
|
|
if t.retryer != nil || t.embedRetryer != nil {
|
|
t.launchEnrichment(created.ID, content)
|
|
}
|
|
|
|
return nil, CaptureOutput{Thought: created}, nil
|
|
}
|
|
|
|
func (t *CaptureTool) launchEnrichment(id uuid.UUID, content string) {
|
|
go func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
|
defer cancel()
|
|
|
|
if t.retryer != nil {
|
|
attemptedAt := time.Now().UTC()
|
|
rawMetadata := metadata.Fallback(t.capture)
|
|
extracted, err := t.provider.ExtractMetadata(ctx, content)
|
|
if err != nil {
|
|
failed := metadata.MarkMetadataFailed(rawMetadata, t.capture, attemptedAt, err)
|
|
if _, updateErr := t.store.UpdateThoughtMetadata(ctx, id, failed); updateErr != nil {
|
|
t.log.Warn("deferred metadata failure could not be persisted",
|
|
slog.String("thought_id", id.String()),
|
|
slog.String("error", updateErr.Error()),
|
|
)
|
|
}
|
|
t.log.Warn("deferred metadata extraction failed",
|
|
slog.String("thought_id", id.String()),
|
|
slog.String("provider", t.provider.Name()),
|
|
slog.String("error", err.Error()),
|
|
)
|
|
t.retryer.QueueThought(id)
|
|
} else {
|
|
completed := metadata.MarkMetadataComplete(extracted, t.capture, attemptedAt)
|
|
if _, updateErr := t.store.UpdateThoughtMetadata(ctx, id, completed); updateErr != nil {
|
|
t.log.Warn("deferred metadata completion could not be persisted",
|
|
slog.String("thought_id", id.String()),
|
|
slog.String("error", updateErr.Error()),
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
if t.embedRetryer != nil {
|
|
if _, err := t.provider.Embed(ctx, content); err != nil {
|
|
t.log.Warn("deferred embedding failed",
|
|
slog.String("thought_id", id.String()),
|
|
slog.String("provider", t.provider.Name()),
|
|
slog.String("error", err.Error()),
|
|
)
|
|
}
|
|
t.embedRetryer.QueueThought(ctx, id, content)
|
|
}
|
|
}()
|
|
}
|