feat: implement 6 tools (search, create, summary, tags, related)

2026-05-05 11:30:49 +08:00
parent d2c5dafabe
commit 18790f1bcd
6 changed files with 262 additions and 0 deletions
@@ -0,0 +1,40 @@
+package tools
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"eino-test/internal/rag"
+	"eino-test/internal/store"
+
+	"github.com/cloudwego/eino/components/tool"
+	"github.com/cloudwego/eino/components/tool/utils"
+)
+
+type CreateNoteInput struct {
+	Title   string   `json:"title" jsonschema:"description=Note title,required"`
+	Content string   `json:"content" jsonschema:"description=Markdown content of the note,required"`
+	Tags    []string `json:"tags" jsonschema:"description=Tags for categorization"`
+}
+
+func NewCreateNoteTool(pipeline *rag.RAGPipeline) tool.InvokableTool {
+	t, _ := utils.InferTool("create_note", "Create a new note in the knowledge base.", func(ctx context.Context, input *CreateNoteInput) (string, error) {
+		note := &store.Note{
+			ID:        input.Title,
+			Title:     input.Title,
+			Content:   input.Content,
+			Tags:      input.Tags,
+			CreatedAt: time.Now(),
+			UpdatedAt: time.Now(),
+		}
+		if err := pipeline.NoteStore().Create(ctx, note); err != nil {
+			return "", fmt.Errorf("create note: %w", err)
+		}
+		if err := pipeline.IndexNote(ctx, note); err != nil {
+			return "", fmt.Errorf("index note: %w", err)
+		}
+		return fmt.Sprintf("Note created: %s", note.ID), nil
+	})
+	return t
+}
@@ -0,0 +1,53 @@
+package tools
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"eino-test/internal/rag"
+
+	"github.com/cloudwego/eino/components/tool"
+	"github.com/cloudwego/eino/components/tool/utils"
+)
+
+type FindRelatedInput struct {
+	NoteID string `json:"note_id" jsonschema:"description=The ID of the note to find related notes for,required"`
+}
+
+func NewFindRelatedTool(pipeline *rag.RAGPipeline) tool.InvokableTool {
+	t, _ := utils.InferTool("find_related", "Find notes related to a given note based on content similarity.", func(ctx context.Context, input *FindRelatedInput) (string, error) {
+		note, err := pipeline.NoteStore().GetByID(ctx, input.NoteID)
+		if err != nil {
+			return "", fmt.Errorf("note not found: %w", err)
+		}
+
+		results, err := pipeline.Search(ctx, note.Content, 6)
+		if err != nil {
+			return "", fmt.Errorf("search related: %w", err)
+		}
+
+		type related struct {
+			Title string  `json:"title"`
+			Score float64 `json:"score"`
+			Tags  []string `json:"tags"`
+		}
+		var items []related
+		for _, r := range results {
+			if r.Note.ID == input.NoteID {
+				continue
+			}
+			items = append(items, related{
+				Title: r.Note.Title,
+				Score: r.Score,
+				Tags:  r.Note.Tags,
+			})
+		}
+		if len(items) > 5 {
+			items = items[:5]
+		}
+		b, _ := json.Marshal(items)
+		return string(b), nil
+	})
+	return t
+}
@@ -0,0 +1,36 @@
+package tools
+
+import (
+	"context"
+	"fmt"
+
+	"eino-test/internal/rag"
+
+	"github.com/cloudwego/eino/components/model"
+	"github.com/cloudwego/eino/components/tool"
+	"github.com/cloudwego/eino/components/tool/utils"
+	"github.com/cloudwego/eino/schema"
+)
+
+type GenerateSummaryInput struct {
+	NoteID string `json:"note_id" jsonschema:"description=The ID of the note to summarize,required"`
+}
+
+func NewGenerateSummaryTool(pipeline *rag.RAGPipeline, chatModel model.BaseChatModel) tool.InvokableTool {
+	t, _ := utils.InferTool("generate_summary", "Generate a concise summary of a note.", func(ctx context.Context, input *GenerateSummaryInput) (string, error) {
+		note, err := pipeline.NoteStore().GetByID(ctx, input.NoteID)
+		if err != nil {
+			return "", fmt.Errorf("note not found: %w", err)
+		}
+		messages := []*schema.Message{
+			schema.SystemMessage("You are a helpful assistant. Generate a concise summary of the following note in 2-3 sentences."),
+			schema.UserMessage(fmt.Sprintf("Title: %s\n\nContent:\n%s", note.Title, note.Content)),
+		}
+		resp, err := chatModel.Generate(ctx, messages)
+		if err != nil {
+			return "", fmt.Errorf("generate summary: %w", err)
+		}
+		return resp.Content, nil
+	})
+	return t
+}
@@ -0,0 +1,45 @@
+package tools
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"eino-test/internal/rag"
+
+	"github.com/cloudwego/eino/components/tool"
+	"github.com/cloudwego/eino/components/tool/utils"
+)
+
+type KeywordSearchInput struct {
+	Query string `json:"query" jsonschema:"description=Keywords to search for,required"`
+}
+
+func NewKeywordSearchTool(pipeline *rag.RAGPipeline) tool.InvokableTool {
+	t, _ := utils.InferTool("keyword_search", "Search notes by exact keyword matching in title and content.", func(ctx context.Context, input *KeywordSearchInput) (string, error) {
+		results, err := pipeline.NoteStore().SearchByKeyword(ctx, input.Query)
+		if err != nil {
+			return "", fmt.Errorf("keyword search failed: %w", err)
+		}
+		type item struct {
+			Title   string   `json:"title"`
+			Content string   `json:"content_preview"`
+			Tags    []string `json:"tags"`
+		}
+		items := make([]item, 0, len(results))
+		for _, n := range results {
+			preview := n.Content
+			if len(preview) > 200 {
+				preview = preview[:200] + "..."
+			}
+			items = append(items, item{
+				Title:   n.Title,
+				Content: preview,
+				Tags:    n.Tags,
+			})
+		}
+		b, _ := json.Marshal(items)
+		return string(b), nil
+	})
+	return t
+}
@@ -0,0 +1,52 @@
+package tools
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"eino-test/internal/rag"
+
+	"github.com/cloudwego/eino/components/tool"
+	"github.com/cloudwego/eino/components/tool/utils"
+)
+
+type SemanticSearchInput struct {
+	Query string `json:"query" jsonschema:"description=The search query,required"`
+	TopK  int    `json:"top_k" jsonschema:"description=Number of results to return"`
+}
+
+func NewSemanticSearchTool(pipeline *rag.RAGPipeline) tool.InvokableTool {
+	t, _ := utils.InferTool("semantic_search", "Search notes by semantic similarity. Use this for conceptual or fuzzy queries.", func(ctx context.Context, input *SemanticSearchInput) (string, error) {
+		topK := input.TopK
+		if topK <= 0 {
+			topK = 5
+		}
+		results, err := pipeline.Search(ctx, input.Query, topK)
+		if err != nil {
+			return "", fmt.Errorf("search failed: %w", err)
+		}
+		type item struct {
+			Title   string   `json:"title"`
+			Content string   `json:"content_preview"`
+			Score   float64  `json:"score"`
+			Tags    []string `json:"tags"`
+		}
+		items := make([]item, 0, len(results))
+		for _, r := range results {
+			preview := r.Note.Content
+			if len(preview) > 200 {
+				preview = preview[:200] + "..."
+			}
+			items = append(items, item{
+				Title:   r.Note.Title,
+				Content: preview,
+				Score:   r.Score,
+				Tags:    r.Note.Tags,
+			})
+		}
+		b, _ := json.Marshal(items)
+		return string(b), nil
+	})
+	return t
+}
@@ -0,0 +1,36 @@
+package tools
+
+import (
+	"context"
+	"fmt"
+
+	"eino-test/internal/rag"
+
+	"github.com/cloudwego/eino/components/model"
+	"github.com/cloudwego/eino/components/tool"
+	"github.com/cloudwego/eino/components/tool/utils"
+	"github.com/cloudwego/eino/schema"
+)
+
+type SuggestTagsInput struct {
+	NoteID string `json:"note_id" jsonschema:"description=The ID of the note to suggest tags for,required"`
+}
+
+func NewSuggestTagsTool(pipeline *rag.RAGPipeline, chatModel model.BaseChatModel) tool.InvokableTool {
+	t, _ := utils.InferTool("suggest_tags", "Suggest relevant tags for a note based on its content.", func(ctx context.Context, input *SuggestTagsInput) (string, error) {
+		note, err := pipeline.NoteStore().GetByID(ctx, input.NoteID)
+		if err != nil {
+			return "", fmt.Errorf("note not found: %w", err)
+		}
+		messages := []*schema.Message{
+			schema.SystemMessage("You are a helpful assistant. Suggest 3-5 relevant tags for the following note. Return only a JSON array of tag strings, e.g. [\"tag1\", \"tag2\"]."),
+			schema.UserMessage(fmt.Sprintf("Title: %s\n\nContent:\n%s", note.Title, note.Content)),
+		}
+		resp, err := chatModel.Generate(ctx, messages)
+		if err != nil {
+			return "", fmt.Errorf("suggest tags: %w", err)
+		}
+		return resp.Content, nil
+	})
+	return t
+}