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