Hein
14e218d784
Some checks failed
CI / build-and-test (push) Failing after -32m22s
* Implement tests for migrating configuration from v1 to v2 for the litellm provider. * Validate the structure and values of the migrated configuration. * Ensure migration rejects newer versions of the configuration. fix(validate): enhance AI provider validation logic * Consolidate provider validation into a dedicated method. * Ensure at least one provider is specified and validate its type. * Check for required fields based on provider type. fix(mcpserver): update tool set to use new enrichment tool * Replace RetryMetadataTool with RetryEnrichmentTool in the ToolSet. fix(tools): refactor tools to use embedding and metadata runners * Update tools to utilize EmbeddingRunner and MetadataRunner instead of Provider. * Adjust method calls to align with the new runner interfaces.
96 lines
3.1 KiB
Go
96 lines
3.1 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
"strings"
|
|
"time"
|
|
|
|
"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/store"
|
|
thoughttypes "git.warky.dev/wdevs/amcs/internal/types"
|
|
)
|
|
|
|
type UpdateTool struct {
|
|
store *store.DB
|
|
embeddings *ai.EmbeddingRunner
|
|
metadata *ai.MetadataRunner
|
|
capture config.CaptureConfig
|
|
log *slog.Logger
|
|
}
|
|
|
|
type UpdateInput struct {
|
|
ID string `json:"id" jsonschema:"the thought id"`
|
|
Content *string `json:"content,omitempty" jsonschema:"replacement content for the thought"`
|
|
Metadata thoughttypes.ThoughtMetadata `json:"metadata,omitempty" jsonschema:"metadata fields to merge into the thought"`
|
|
Project string `json:"project,omitempty" jsonschema:"optional project name or id to move the thought into"`
|
|
}
|
|
|
|
type UpdateOutput struct {
|
|
Thought thoughttypes.Thought `json:"thought"`
|
|
}
|
|
|
|
func NewUpdateTool(db *store.DB, embeddings *ai.EmbeddingRunner, metadata *ai.MetadataRunner, capture config.CaptureConfig, log *slog.Logger) *UpdateTool {
|
|
return &UpdateTool{store: db, embeddings: embeddings, metadata: metadata, capture: capture, log: log}
|
|
}
|
|
|
|
func (t *UpdateTool) Handle(ctx context.Context, _ *mcp.CallToolRequest, in UpdateInput) (*mcp.CallToolResult, UpdateOutput, error) {
|
|
id, err := parseUUID(in.ID)
|
|
if err != nil {
|
|
return nil, UpdateOutput{}, err
|
|
}
|
|
|
|
current, err := t.store.GetThought(ctx, id)
|
|
if err != nil {
|
|
return nil, UpdateOutput{}, err
|
|
}
|
|
|
|
content := current.Content
|
|
var embedding []float32
|
|
embeddingModel := ""
|
|
mergedMetadata := current.Metadata
|
|
projectID := current.ProjectID
|
|
|
|
if in.Content != nil {
|
|
content = strings.TrimSpace(*in.Content)
|
|
if content == "" {
|
|
return nil, UpdateOutput{}, errInvalidInput("content must not be empty")
|
|
}
|
|
embedResult, err := t.embeddings.Embed(ctx, content)
|
|
if err != nil {
|
|
return nil, UpdateOutput{}, err
|
|
}
|
|
embedding = embedResult.Vector
|
|
embeddingModel = embedResult.Model
|
|
extracted, extractErr := t.metadata.ExtractMetadata(ctx, content)
|
|
if extractErr != nil {
|
|
t.log.Warn("metadata extraction failed during update, keeping current metadata", slog.String("error", extractErr.Error()))
|
|
mergedMetadata = metadata.MarkMetadataFailed(mergedMetadata, t.capture, time.Now().UTC(), extractErr)
|
|
} else {
|
|
mergedMetadata = metadata.MarkMetadataComplete(metadata.SanitizeExtracted(extracted), t.capture, time.Now().UTC())
|
|
mergedMetadata.Attachments = current.Metadata.Attachments
|
|
}
|
|
}
|
|
|
|
mergedMetadata = metadata.Merge(mergedMetadata, in.Metadata, t.capture)
|
|
|
|
if rawProject := strings.TrimSpace(in.Project); rawProject != "" {
|
|
project, err := t.store.GetProject(ctx, rawProject)
|
|
if err != nil {
|
|
return nil, UpdateOutput{}, err
|
|
}
|
|
projectID = &project.ID
|
|
}
|
|
|
|
updated, err := t.store.UpdateThought(ctx, id, content, embedding, embeddingModel, mergedMetadata, projectID)
|
|
if err != nil {
|
|
return nil, UpdateOutput{}, err
|
|
}
|
|
|
|
return nil, UpdateOutput{Thought: updated}, nil
|
|
}
|