diff --git a/internal/tools/create_note.go b/internal/tools/create_note.go new file mode 100644 index 0000000..ac7af43 --- /dev/null +++ b/internal/tools/create_note.go @@ -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 +} diff --git a/internal/tools/find_related.go b/internal/tools/find_related.go new file mode 100644 index 0000000..8a83934 --- /dev/null +++ b/internal/tools/find_related.go @@ -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 +} diff --git a/internal/tools/generate_summary.go b/internal/tools/generate_summary.go new file mode 100644 index 0000000..8f1bd99 --- /dev/null +++ b/internal/tools/generate_summary.go @@ -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 +} diff --git a/internal/tools/keyword_search.go b/internal/tools/keyword_search.go new file mode 100644 index 0000000..71b3852 --- /dev/null +++ b/internal/tools/keyword_search.go @@ -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 +} diff --git a/internal/tools/semantic_search.go b/internal/tools/semantic_search.go new file mode 100644 index 0000000..3f4a41c --- /dev/null +++ b/internal/tools/semantic_search.go @@ -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 +} diff --git a/internal/tools/suggest_tags.go b/internal/tools/suggest_tags.go new file mode 100644 index 0000000..599fef6 --- /dev/null +++ b/internal/tools/suggest_tags.go @@ -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 +}