53 lines
1.4 KiB
Go
53 lines
1.4 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 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
|
|
}
|