forked from tmc/langchaingo
-
Notifications
You must be signed in to change notification settings - Fork 0
/
vectorstores.go
53 lines (43 loc) · 1.43 KB
/
vectorstores.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
package vectorstores
import (
"context"
"github.com/portyl/langchaingo/callbacks"
"github.com/portyl/langchaingo/schema"
)
// VectorStore is the interface for saving and querying documents in the
// form of vector embeddings.
type VectorStore interface {
AddDocuments(context.Context, []schema.Document, ...Option) error
SimilaritySearch(ctx context.Context, query string, numDocuments int, options ...Option) ([]schema.Document, error) //nolint:lll
}
// Retriever is a retriever for vector stores.
type Retriever struct {
CallbacksHandler callbacks.Handler
v VectorStore
numDocs int
options []Option
}
var _ schema.Retriever = Retriever{}
// GetRelevantDocuments returns documents using the vector store.
func (r Retriever) GetRelevantDocuments(ctx context.Context, query string) ([]schema.Document, error) {
if r.CallbacksHandler != nil {
r.CallbacksHandler.HandleRetrieverStart(ctx, query)
}
docs, err := r.v.SimilaritySearch(ctx, query, r.numDocs, r.options...)
if err != nil {
return nil, err
}
if r.CallbacksHandler != nil {
r.CallbacksHandler.HandleRetrieverEnd(ctx, docs)
}
return docs, nil
}
// ToRetriever takes a vector store and returns a retriever using the
// vector store to retrieve documents.
func ToRetriever(vectorStore VectorStore, numDocuments int, options ...Option) Retriever {
return Retriever{
v: vectorStore,
numDocs: numDocuments,
options: options,
}
}