chore: remove retriever (#132)

henomis · Sep 14, 2023 · 7c6d090 · 7c6d090
1 parent 53a2af9
commit 7c6d090
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 79 deletions.
diff --git a/README.md b/README.md
@@ -58,7 +58,7 @@ import (
 	"context"
 
 	openaiembedder "github.com/henomis/lingoose/embedder/openai"
-	"github.com/henomis/lingoose/index/retriever"
+	"github.com/henomis/lingoose/index/option"
 	simplevectorindex "github.com/henomis/lingoose/index/simpleVectorIndex"
 	"github.com/henomis/lingoose/llm/openai"
 	"github.com/henomis/lingoose/loader"
@@ -67,12 +67,14 @@ import (
 )
 
 func main() {
-	docs, _ := loader.NewPDFToTextLoader("./kb").WithTextSplitter(textsplitter.NewRecursiveCharacterTextSplitter(2000, 200)).Load(context.Background())
-	qapipeline.New(openai.NewChat().WithVerbose(true)).WithRetriever(retriever.New(simplevectorindex.New("db", ".", openaiembedder.New(openaiembedder.AdaEmbeddingV2)), &docs)).Query(context.Background(), "What is the NATO purpose?")
+	docs, _ := loader.NewPDFToTextLoader("./kb").WithPDFToTextPath("/opt/homebrew/bin/pdftotext").WithTextSplitter(textsplitter.NewRecursiveCharacterTextSplitter(2000, 200)).Load(context.Background())
+	index := simplevectorindex.New("db", ".", openaiembedder.New(openaiembedder.AdaEmbeddingV2))
+	index.LoadFromDocuments(context.Background(), docs)
+	qapipeline.New(openai.NewChat().WithVerbose(true)).WithIndex(index).Query(context.Background(), "What is the NATO purpose?", option.WithTopK(1))
 }
 ```
 
-This is the _famous_ 2-lines **lingoose** knowledge base chatbot. 🤖
+This is the _famous_ 4-lines **lingoose** knowledge base chatbot. 🤖
 
 # Installation
 

diff --git a/examples/embeddings/simplekb/main.go b/examples/embeddings/simplekb/main.go
@@ -4,7 +4,7 @@ import (
 	"context"
 
 	openaiembedder "github.com/henomis/lingoose/embedder/openai"
-	"github.com/henomis/lingoose/index/retriever"
+	"github.com/henomis/lingoose/index/option"
 	simplevectorindex "github.com/henomis/lingoose/index/simpleVectorIndex"
 	"github.com/henomis/lingoose/llm/openai"
 	"github.com/henomis/lingoose/loader"
@@ -13,6 +13,8 @@ import (
 )
 
 func main() {
-	docs, _ := loader.NewPDFToTextLoader("./kb").WithTextSplitter(textsplitter.NewRecursiveCharacterTextSplitter(2000, 200)).Load(context.Background())
-	qapipeline.New(openai.NewChat().WithVerbose(true)).WithRetriever(retriever.New(simplevectorindex.New("db", ".", openaiembedder.New(openaiembedder.AdaEmbeddingV2)), &docs)).Query(context.Background(), "What is the NATO purpose?")
+	docs, _ := loader.NewPDFToTextLoader("./kb").WithPDFToTextPath("/opt/homebrew/bin/pdftotext").WithTextSplitter(textsplitter.NewRecursiveCharacterTextSplitter(2000, 200)).Load(context.Background())
+	index := simplevectorindex.New("db", ".", openaiembedder.New(openaiembedder.AdaEmbeddingV2))
+	index.LoadFromDocuments(context.Background(), docs)
+	qapipeline.New(openai.NewChat().WithVerbose(true)).WithIndex(index).Query(context.Background(), "What is the NATO purpose?", option.WithTopK(1))
 }
diff --git a/index/retriever/retriever.go b/index/retriever/retriever.go
diff --git a/pipeline/qa/qa.go b/pipeline/qa/qa.go
@@ -6,6 +6,8 @@ import (
 
 	"github.com/henomis/lingoose/chat"
 	"github.com/henomis/lingoose/document"
+	"github.com/henomis/lingoose/index"
+	indexoption "github.com/henomis/lingoose/index/option"
 	"github.com/henomis/lingoose/pipeline"
 	"github.com/henomis/lingoose/prompt"
 	"github.com/henomis/lingoose/types"
@@ -19,14 +21,14 @@ const (
 	qaTubeUserPromptTemplate = "Based on the following context answer to the question.\n\nContext:\n{{.context}}\n\nQuestion: {{.query}}"
 )
 
-type Retriever interface {
-	Query(context.Context, string) ([]document.Document, error)
+type Index interface {
+	Query(context.Context, string, ...indexoption.Option) (index.SearchResults, error)
 }
 
 type QAPipeline struct {
 	llmEngine pipeline.LlmEngine
 	pipeline  *pipeline.Pipeline
-	retriever Retriever
+	index     Index
 }
 
 func New(llmEngine pipeline.LlmEngine) *QAPipeline {
@@ -55,7 +57,7 @@ func New(llmEngine pipeline.LlmEngine) *QAPipeline {
 	return &QAPipeline{
 		llmEngine: llmEngine,
 		pipeline:  pipeline.New(tube),
-		retriever: nil,
+		index:     nil,
 	}
 }
 
@@ -73,22 +75,22 @@ func (p *QAPipeline) WithPrompt(chat *chat.Chat) *QAPipeline {
 	}
 }
 
-func (p *QAPipeline) WithRetriever(retriever Retriever) *QAPipeline {
-	p.retriever = retriever
-	return p
+func (q *QAPipeline) WithIndex(index Index) *QAPipeline {
+	q.index = index
+	return q
 }
 
-func (q *QAPipeline) Query(ctx context.Context, query string) (types.M, error) {
-	if q.retriever == nil {
+func (q *QAPipeline) Query(ctx context.Context, query string, opts ...indexoption.Option) (types.M, error) {
+	if q.index == nil {
 		return nil, fmt.Errorf("retriever is not defined")
 	}
 
-	docs, err := q.retriever.Query(ctx, query)
+	docs, err := q.index.Query(ctx, query, opts...)
 	if err != nil {
 		return nil, err
 	}
 
-	return q.Run(ctx, query, docs)
+	return q.Run(ctx, query, docs.ToDocuments())
 }
 
 func (t *QAPipeline) Run(ctx context.Context, query string, documents []document.Document) (types.M, error) {