118 lines
2.6 KiB
Go
118 lines
2.6 KiB
Go
|
package eino
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"errors"
|
|||
|
|
|
|||
|
|
"github.com/cloudwego/eino/callbacks"
|
|||
|
|
"github.com/cloudwego/eino/components/embedding"
|
|||
|
|
"github.com/cloudwego/eino/components/retriever"
|
|||
|
|
"github.com/cloudwego/eino/schema"
|
|||
|
|
"github.com/gogf/gf/v2/util/gconv"
|
|||
|
|
"github.com/pgvector/pgvector-go"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
type PGVectorRetrieverConfig struct {
|
|||
|
|
Embedder embedding.Embedder
|
|||
|
|
DefaultTopK int
|
|||
|
|
DefaultIndex string
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type PGVectorRetriever struct {
|
|||
|
|
embedder embedding.Embedder
|
|||
|
|
topK int
|
|||
|
|
index string
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func NewPGVectorRetriever(config *PGVectorRetrieverConfig) (*PGVectorRetriever, error) {
|
|||
|
|
if config.Embedder == nil {
|
|||
|
|
return nil, errors.New("embedder is required")
|
|||
|
|
}
|
|||
|
|
if config.DefaultTopK <= 0 {
|
|||
|
|
config.DefaultTopK = 5
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return &PGVectorRetriever{
|
|||
|
|
embedder: config.Embedder,
|
|||
|
|
topK: config.DefaultTopK,
|
|||
|
|
index: config.DefaultIndex,
|
|||
|
|
}, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *PGVectorRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {
|
|||
|
|
|
|||
|
|
// 1. 处理公共 Option（官方标准写法）
|
|||
|
|
options := &retriever.Options{
|
|||
|
|
Index: &r.index,
|
|||
|
|
TopK: &r.topK,
|
|||
|
|
Embedding: r.embedder,
|
|||
|
|
}
|
|||
|
|
options = retriever.GetCommonOptions(options, opts...)
|
|||
|
|
|
|||
|
|
// 2. 回调（官方标准）
|
|||
|
|
ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{
|
|||
|
|
Query: query,
|
|||
|
|
TopK: *options.TopK,
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
// 3. 执行检索
|
|||
|
|
docs, err := r.doRetrieve(ctx, query, options)
|
|||
|
|
if err != nil {
|
|||
|
|
callbacks.OnError(ctx, err)
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 4. 完成回调
|
|||
|
|
callbacks.OnEnd(ctx, &retriever.CallbackOutput{
|
|||
|
|
Docs: docs,
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return docs, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *PGVectorRetriever) doRetrieve(ctx context.Context, query string, opts *retriever.Options) ([]*schema.Document, error) {
|
|||
|
|
|
|||
|
|
// 1. 生成向量
|
|||
|
|
vectors, err := opts.Embedding.EmbedStrings(ctx, []string{query})
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
if len(vectors) == 0 {
|
|||
|
|
return nil, errors.New("empty query vector")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
queryVec := pgvector.NewVector(vectors[0])
|
|||
|
|
topK := *opts.TopK
|
|||
|
|
|
|||
|
|
// 2. PG 向量相似度检索 SQL
|
|||
|
|
sql := `
|
|||
|
|
SELECT id, content, dataset_id, document_id,
|
|||
|
|
vector <-> ? AS distance
|
|||
|
|
FROM document_chunk
|
|||
|
|
ORDER BY distance ASC
|
|||
|
|
LIMIT ?
|
|||
|
|
`
|
|||
|
|
|
|||
|
|
// 3. 查询
|
|||
|
|
rows, err := dao.DocumentChunk.GetDB().GetAll(ctx, sql, queryVec, topK)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 4. 转为 Eino Document
|
|||
|
|
docs := make([]*schema.Document, 0, len(rows))
|
|||
|
|
for _, row := range rows {
|
|||
|
|
docs = append(docs, &schema.Document{
|
|||
|
|
ID: gconv.String(row["id"]),
|
|||
|
|
Content: gconv.String(row["content"]),
|
|||
|
|
Metadata: map[string]any{
|
|||
|
|
"dataset_id": row["dataset_id"],
|
|||
|
|
"document_id": row["document_id"],
|
|||
|
|
"distance": row["distance"],
|
|||
|
|
},
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return docs, nil
|
|||
|
|
}
|