amcs/internal/tools/links.go

package tools

import (
	"context"
	"strings"

	"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/store"
	thoughttypes "git.warky.dev/wdevs/amcs/internal/types"
)

type LinksTool struct {
	store    *store.DB
	provider ai.Provider
	search   config.SearchConfig
}

type LinkInput struct {
	FromID   string `json:"from_id" jsonschema:"the source thought id"`
	ToID     string `json:"to_id" jsonschema:"the target thought id"`
	Relation string `json:"relation" jsonschema:"relationship label such as follows_up or references"`
}

type LinkOutput struct {
	Linked bool `json:"linked"`
}

type RelatedInput struct {
	ID              string `json:"id" jsonschema:"the thought id"`
	IncludeSemantic *bool  `json:"include_semantic,omitempty" jsonschema:"whether to include semantic neighbors; defaults to true"`
}

type RelatedThought struct {
	ID         string                       `json:"id"`
	Content    string                       `json:"content"`
	Metadata   thoughttypes.ThoughtMetadata `json:"metadata"`
	Relation   string                       `json:"relation,omitempty"`
	Direction  string                       `json:"direction,omitempty"`
	Similarity float64                      `json:"similarity,omitempty"`
	Source     string                       `json:"source"`
}

type RelatedOutput struct {
	Related []RelatedThought `json:"related"`
}

func NewLinksTool(db *store.DB, provider ai.Provider, search config.SearchConfig) *LinksTool {
	return &LinksTool{store: db, provider: provider, search: search}
}

func (t *LinksTool) Link(ctx context.Context, _ *mcp.CallToolRequest, in LinkInput) (*mcp.CallToolResult, LinkOutput, error) {
	fromID, err := parseUUID(in.FromID)
	if err != nil {
		return nil, LinkOutput{}, err
	}
	toID, err := parseUUID(in.ToID)
	if err != nil {
		return nil, LinkOutput{}, err
	}
	relation := strings.TrimSpace(in.Relation)
	if relation == "" {
		return nil, LinkOutput{}, errInvalidInput("relation is required")
	}
	if _, err := t.store.GetThought(ctx, fromID); err != nil {
		return nil, LinkOutput{}, err
	}
	if _, err := t.store.GetThought(ctx, toID); err != nil {
		return nil, LinkOutput{}, err
	}
	if err := t.store.InsertLink(ctx, thoughttypes.ThoughtLink{
		FromID:   fromID,
		ToID:     toID,
		Relation: relation,
	}); err != nil {
		return nil, LinkOutput{}, err
	}
	return nil, LinkOutput{Linked: true}, nil
}

func (t *LinksTool) Related(ctx context.Context, _ *mcp.CallToolRequest, in RelatedInput) (*mcp.CallToolResult, RelatedOutput, error) {
	id, err := parseUUID(in.ID)
	if err != nil {
		return nil, RelatedOutput{}, err
	}

	thought, err := t.store.GetThought(ctx, id)
	if err != nil {
		return nil, RelatedOutput{}, err
	}

	linked, err := t.store.LinkedThoughts(ctx, id)
	if err != nil {
		return nil, RelatedOutput{}, err
	}

	related := make([]RelatedThought, 0, len(linked)+t.search.DefaultLimit)
	seen := map[string]struct{}{thought.ID.String(): {}}
	for _, item := range linked {
		key := item.Thought.ID.String()
		seen[key] = struct{}{}
		related = append(related, RelatedThought{
			ID:        key,
			Content:   item.Thought.Content,
			Metadata:  item.Thought.Metadata,
			Relation:  item.Relation,
			Direction: item.Direction,
			Source:    "link",
		})
	}

	includeSemantic := true
	if in.IncludeSemantic != nil {
		includeSemantic = *in.IncludeSemantic
	}

	if includeSemantic {
		embedding, err := t.provider.Embed(ctx, thought.Content)
		if err != nil {
			return nil, RelatedOutput{}, err
		}
		semantic, err := t.store.SearchSimilarThoughts(ctx, embedding, t.provider.EmbeddingModel(), t.search.DefaultThreshold, t.search.DefaultLimit, thought.ProjectID, &thought.ID)
		if err != nil {
			return nil, RelatedOutput{}, err
		}
		for _, item := range semantic {
			key := item.ID.String()
			if _, ok := seen[key]; ok {
				continue
			}
			seen[key] = struct{}{}
			related = append(related, RelatedThought{
				ID:         key,
				Content:    item.Content,
				Metadata:   item.Metadata,
				Similarity: item.Similarity,
				Source:     "semantic",
			})
		}
	}

	return nil, RelatedOutput{Related: related}, nil
}