37 lines
1.2 KiB
Go
37 lines
1.2 KiB
Go
|
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
|
||
|
|
}
|