54 lines
1.3 KiB
Go
54 lines
1.3 KiB
Go
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
|
|
}
|