Skip to content

Commit

Permalink
Feat Add redis as vector database (#145)
Browse files Browse the repository at this point in the history
  • Loading branch information
henomis authored Nov 2, 2023
1 parent e8af1a4 commit 83195d2
Show file tree
Hide file tree
Showing 8 changed files with 372 additions and 17 deletions.
7 changes: 6 additions & 1 deletion .github/workflows/checks.yml
Original file line number Diff line number Diff line change
Expand Up @@ -33,4 +33,9 @@ jobs:
cache: false
- uses: actions/checkout@v3
- name: golangci-lint
uses: golangci/golangci-lint-action@v3
uses: golangci/golangci-lint-action@v3
with:
# Require: The version of golangci-lint to use.
# When `install-mode` is `binary` (default) the value can be v1.2 or v1.2.3 or `latest` to use the latest version.
# When `install-mode` is `goinstall` the value can be v1.2.3, `latest`, or the hash of a commit.
version: v1.54.2
1 change: 0 additions & 1 deletion examples/embeddings/qdrant/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ func main() {
qdrantdb.New(
qdrantdb.Options{
CollectionName: "test",
IncludeContent: true,
CreateCollection: &qdrantdb.CreateCollectionOptions{
Dimension: 1536,
Distance: qdrantdb.DistanceCosine,
Expand Down
107 changes: 107 additions & 0 deletions examples/embeddings/redis/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
package main

import (
"context"
"fmt"

"github.com/RediSearch/redisearch-go/v2/redisearch"
openaiembedder "github.com/henomis/lingoose/embedder/openai"
"github.com/henomis/lingoose/index"
indexoption "github.com/henomis/lingoose/index/option"
"github.com/henomis/lingoose/index/vectordb/redis"
"github.com/henomis/lingoose/llm/openai"
"github.com/henomis/lingoose/loader"
"github.com/henomis/lingoose/prompt"
"github.com/henomis/lingoose/textsplitter"
)

// download https://raw.githubusercontent.com/hwchase17/chat-your-data/master/state_of_the_union.txt

func main() {
index := index.New(
redis.New(
redis.Options{
RedisearchClient: redisearch.NewClient("localhost:6379", "test"),
CreateIndex: &redis.CreateIndexOptions{
Dimension: 1536,
Distance: redis.DistanceCosine,
},
},
),
openaiembedder.New(openaiembedder.AdaEmbeddingV2),
).WithIncludeContents(true)

indexIsEmpty, err := index.IsEmpty(context.Background())
if err != nil {
panic(err)
}

if indexIsEmpty {
err = ingestData(index)
if err != nil {
panic(err)
}
}

query := "What is the purpose of the NATO Alliance?"
similarities, err := index.Query(
context.Background(),
query,
indexoption.WithTopK(3),
)
if err != nil {
panic(err)
}

content := ""
for _, similarity := range similarities {
fmt.Printf("Similarity: %f\n", similarity.Score)
fmt.Printf("Document: %s\n", similarity.Content())
fmt.Println("Metadata: ", similarity.Metadata)
fmt.Println("ID: ", similarity.ID)
fmt.Println("----------")
content += similarity.Content() + "\n"
}

llmOpenAI := openai.NewCompletion().WithVerbose(true)

prompt1 := prompt.NewPromptTemplate(
"Based on the following context answer to the question.\n\nContext:\n{{.context}}\n\nQuestion: {{.query}}").
WithInputs(
map[string]string{
"query": query,
"context": content,
},
)

err = prompt1.Format(nil)
if err != nil {
panic(err)
}

_, err = llmOpenAI.Completion(context.Background(), prompt1.String())
if err != nil {
panic(err)
}
}

func ingestData(redisIndex *index.Index) error {
documents, err := loader.NewDirectoryLoader(".", ".txt").Load(context.Background())
if err != nil {
return err
}

textSplitter := textsplitter.NewRecursiveCharacterTextSplitter(1000, 20)

documentChunks := textSplitter.SplitDocuments(documents)

for _, doc := range documentChunks {
fmt.Println(doc.Content)
fmt.Println("----------")
fmt.Println(doc.Metadata)
fmt.Println("----------")
fmt.Println()
}

return redisIndex.LoadFromDocuments(context.Background(), documentChunks)
}
2 changes: 2 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ toolchain go1.21.1
require github.com/mitchellh/mapstructure v1.5.0

require (
github.com/RediSearch/redisearch-go/v2 v2.1.1
github.com/google/uuid v1.3.0
github.com/henomis/cohere-go v1.0.1
github.com/henomis/milvus-go v0.0.4
Expand All @@ -19,6 +20,7 @@ require (

require (
github.com/dlclark/regexp2 v1.8.1 // indirect
github.com/gomodule/redigo v1.8.9 // indirect
github.com/henomis/restclientgo v1.0.6 // indirect
github.com/iancoleman/orderedmap v0.0.0-20190318233801-ac98e3ecb4b0 // indirect
)
10 changes: 10 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
github.com/RediSearch/redisearch-go/v2 v2.1.1 h1:cCn3i40uLsVD8cxwrdrGfhdAgbR5Cld9q11eYyVOwpM=
github.com/RediSearch/redisearch-go/v2 v2.1.1/go.mod h1:Uw93Wi97QqAsw1DwbQrhVd88dBorGTfSuCS42zfh1iA=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0=
github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/gomodule/redigo v1.8.9 h1:Sl3u+2BI/kk+VEatbj0scLdrFhjPmbxOc1myhDP41ws=
github.com/gomodule/redigo v1.8.9/go.mod h1:7ArFNvsTjH8GMMzB4uy1snslv2BwmginuMs06a1uzZE=
github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/henomis/cohere-go v1.0.1 h1:a47gIN29tqAl4yBTAT+BzQMjsWG94Fz07u9AE4Md+a8=
Expand All @@ -28,8 +32,14 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
github.com/sashabaranov/go-openai v1.12.0 h1:aRNHH0gtVfrpIaEolD0sWrLLRnYQNK4cH/bIAHwL8Rk=
github.com/sashabaranov/go-openai v1.12.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/testify v1.3.1-0.20190311161405-34c6fa2dc709/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
2 changes: 0 additions & 2 deletions index/vectordb/milvus/milvus.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,6 @@ type CreateCollectionOptions struct {
type Options struct {
DatabaseName *string
CollectionName string
IncludeContent bool
IncludeValues bool
BatchUpsertSize *int
CreateCollection *CreateCollectionOptions
}
Expand Down
16 changes: 3 additions & 13 deletions index/vectordb/qdrant/qdrant.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@ import (
type DB struct {
qdrantClient *qdrantgo.Client
collectionName string
includeContent bool
includeValues bool

createCollection *CreateCollectionOptions
}
Expand All @@ -38,9 +36,6 @@ type CreateCollectionOptions struct {

type Options struct {
CollectionName string
IncludeContent bool
IncludeValues bool
BatchUpsertSize *int
CreateCollection *CreateCollectionOptions
}

Expand All @@ -53,8 +48,6 @@ func New(options Options) *DB {
return &DB{
qdrantClient: qdrantClient,
collectionName: options.CollectionName,
includeContent: options.IncludeContent,
includeValues: options.IncludeValues,
createCollection: options.CreateCollection,
}
}
Expand Down Expand Up @@ -126,7 +119,7 @@ func (d *DB) Search(ctx context.Context, values []float64, options *option.Optio
return nil, fmt.Errorf("%w: %w", index.ErrInternal, err)
}

return buildSearchResultsFromQdrantMatches(matches, d.includeContent), nil
return buildSearchResultsFromQdrantMatches(matches), nil
}

func (d *DB) similaritySearch(
Expand All @@ -139,6 +132,7 @@ func (d *DB) similaritySearch(
}

includeMetadata := true
includeValues := true
res := &qdrantresponse.PointSearch{}
err := d.qdrantClient.PointSearch(
ctx,
Expand All @@ -147,7 +141,7 @@ func (d *DB) similaritySearch(
Limit: opts.TopK,
Vector: values,
WithPayload: &includeMetadata,
WithVector: &d.includeValues,
WithVector: &includeValues,
Filter: opts.Filter.(qdrantrequest.Filter),
},
res,
Expand Down Expand Up @@ -195,15 +189,11 @@ func (d *DB) createCollectionIfRequired(ctx context.Context) error {

func buildSearchResultsFromQdrantMatches(
matches []qdrantresponse.PointSearchResult,
includeContent bool,
) index.SearchResults {
searchResults := make([]index.SearchResult, len(matches))

for i, match := range matches {
metadata := index.DeepCopyMetadata(match.Payload)
if !includeContent {
delete(metadata, index.DefaultKeyContent)
}

searchResults[i] = index.SearchResult{
Data: index.Data{
Expand Down
Loading

0 comments on commit 83195d2

Please sign in to comment.