diff --git a/.gitignore b/.gitignore index af737e80d..ea636dbe3 100644 --- a/.gitignore +++ b/.gitignore @@ -42,6 +42,9 @@ output/* .claude/ .claude-trace +# Kiro files +.kiro/ + # Data files *.jsonl *.csv diff --git a/components/indexer/valkey/README.md b/components/indexer/valkey/README.md new file mode 100644 index 000000000..47c549eef --- /dev/null +++ b/components/indexer/valkey/README.md @@ -0,0 +1,147 @@ +# Valkey Indexer + +A Valkey indexer implementation for [Eino](https://github.com/cloudwego/eino) that implements the `Indexer` interface. This component uses Valkey Hashes to store documents with vector embeddings via the [valkey-glide](https://github.com/valkey-io/valkey-glide) client, enabling vector similarity search capabilities. + +## Features + +- Implements `github.com/cloudwego/eino/components/indexer.Indexer` +- Uses Valkey Hashes (HSET) or JSON (JSON.SET) for document storage +- Pipeline batch execution for high throughput +- Automatic embedding generation with configurable batch size +- Custom field mapping via `DocumentToHashes` or `DocumentToJSON` function +- Configurable key prefix for index partitioning +- Full Eino callback integration (OnStart, OnEnd, OnError) + +## Installation + +```bash +go get github.com/cloudwego/eino-ext/components/indexer/valkey@latest +``` + +## Prerequisites + +- Valkey 9.1+ with the Search module (e.g., `valkey/valkey-bundle`) +- [valkey-glide](https://github.com/valkey-io/valkey-glide) Go client (requires CGO and Rust core library) + +## Quick Start + +```go +import ( + "context" + "fmt" + + glide "github.com/valkey-io/valkey-glide/go/v2" + "github.com/valkey-io/valkey-glide/go/v2/config" + "github.com/cloudwego/eino/schema" + + valkeyIndexer "github.com/cloudwego/eino-ext/components/indexer/valkey" +) + +func main() { + ctx := context.Background() + + // 1. Create Valkey GLIDE client + cfg := config.NewClientConfiguration(). + WithAddress(&config.NodeAddress{Host: "localhost", Port: 6379}) + client, _ := glide.NewClient(cfg) + + // 2. Create embedding component (use your preferred embedder) + emb := yourEmbedder() + + // 3. Create Valkey indexer + indexer, _ := valkeyIndexer.NewIndexer(ctx, &valkeyIndexer.IndexerConfig{ + Client: client, + KeyPrefix: "doc:", + BatchSize: 10, + Embedding: emb, + }) + + // 4. Store documents + docs := []*schema.Document{ + {ID: "1", Content: "Valkey is a high-performance key-value store."}, + {ID: "2", Content: "Vector search enables semantic similarity queries."}, + } + + ids, err := indexer.Store(ctx, docs) + if err != nil { + fmt.Printf("store error: %v\n", err) + return + } + fmt.Printf("stored document IDs: %v\n", ids) +} +``` + +## Configuration + +```go +type IndexerConfig struct { + // Required: Valkey GLIDE client (must implement Exec for pipeline batching) + Client BatchClient + + // Optional: Key prefix prepended to each hash key + // Should match the prefix used in FT.CREATE + KeyPrefix string + + // Optional: Storage format - DocumentTypeHash (default) or DocumentTypeJSON + DocumentType DocumentType + + // Optional: Custom document-to-hash conversion (Hash mode only) + // Default: stores Content with embedding, plus all MetaData fields + DocumentToHashes func(ctx context.Context, doc *schema.Document) (*Hashes, error) + + // Optional: Custom document-to-JSON conversion (JSON mode only) + // Default: stores content, vector, and metadata fields + DocumentToJSON func(ctx context.Context, doc *schema.Document, vector []float64) (map[string]any, error) + + // Optional: Max texts per embedding batch call (default: 10) + BatchSize int + + // Required: Embedding method for vectorizing document content + Embedding embedding.Embedder +} +``` + +## JSON Document Storage + +For JSON storage, use `DocumentTypeJSON`. This requires the Valkey JSON module and an index created with `ON JSON`: + +```go +indexer, _ := valkeyIndexer.NewIndexer(ctx, &valkeyIndexer.IndexerConfig{ + Client: client, + KeyPrefix: "jdoc:", + DocumentType: valkeyIndexer.DocumentTypeJSON, + Embedding: emb, +}) +``` + +Create the index with JSONPath field identifiers: +``` +FT.CREATE my_json_index ON JSON PREFIX 1 jdoc: SCHEMA + $.content AS content TEXT + $.vector_content AS vector_content VECTOR HNSW 6 TYPE FLOAT32 DIM 1024 DISTANCE_METRIC COSINE +``` + +## Custom Field Mapping + +```go +indexer, _ := valkeyIndexer.NewIndexer(ctx, &valkeyIndexer.IndexerConfig{ + Client: client, + KeyPrefix: "doc:", + DocumentToHashes: func(ctx context.Context, doc *schema.Document) (*valkeyIndexer.Hashes, error) { + return &valkeyIndexer.Hashes{ + Key: doc.ID, + Field2Value: map[string]valkeyIndexer.FieldValue{ + "title": {Value: doc.MetaData["title"]}, + "body": {Value: doc.Content, EmbedKey: "body_vector"}, + }, + }, nil + }, + Embedding: emb, +}) +``` + +## For More Details + +- [Eino Documentation](https://www.cloudwego.io/zh/docs/eino/) +- [Valkey GLIDE Go Client](https://github.com/valkey-io/valkey-glide) +- [Valkey Search](https://valkey.io/topics/search/) diff --git a/components/indexer/valkey/consts.go b/components/indexer/valkey/consts.go new file mode 100644 index 000000000..9dbcaf363 --- /dev/null +++ b/components/indexer/valkey/consts.go @@ -0,0 +1,22 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package valkey + +const ( + defaultReturnFieldContent = "content" + defaultReturnFieldVectorContent = "vector_content" +) diff --git a/components/indexer/valkey/examples/hash_indexer/hash_indexer.go b/components/indexer/valkey/examples/hash_indexer/hash_indexer.go new file mode 100644 index 000000000..ac0b0ec46 --- /dev/null +++ b/components/indexer/valkey/examples/hash_indexer/hash_indexer.go @@ -0,0 +1,90 @@ +//go:build ignore + +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package main + +import ( + "context" + "fmt" + + glide "github.com/valkey-io/valkey-glide/go/v2" + "github.com/valkey-io/valkey-glide/go/v2/config" + + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/schema" + + vi "github.com/cloudwego/eino-ext/components/indexer/valkey" +) + +// This example demonstrates indexing documents as Valkey Hashes. +// Prerequisites: +// - Valkey 9.1+ with Search module running on localhost:6379 +// - Create an index first: +// FT.CREATE my_index ON HASH PREFIX 1 doc: SCHEMA +// content TEXT vector_content VECTOR HNSW 6 TYPE FLOAT32 DIM 1024 DISTANCE_METRIC COSINE +func main() { + ctx := context.Background() + + // 1. Create Valkey GLIDE client + cfg := config.NewClientConfiguration(). + WithAddress(&config.NodeAddress{Host: "localhost", Port: 6379}) + client, err := glide.NewClient(cfg) + if err != nil { + panic(fmt.Sprintf("failed to create client: %v", err)) + } + defer client.Close() + + // 2. Create indexer + indexer, err := vi.NewIndexer(ctx, &vi.IndexerConfig{ + Client: client, + KeyPrefix: "doc:", + DocumentType: vi.DocumentTypeHash, + BatchSize: 10, + Embedding: &mockEmbedding{dims: 1024}, // replace with real embedder + }) + if err != nil { + panic(err) + } + + // 3. Store documents + docs := []*schema.Document{ + {ID: "1", Content: "Eiffel Tower: Located in Paris, France"}, + {ID: "2", Content: "Great Wall of China: One of the greatest wonders of the world"}, + {ID: "3", Content: "Grand Canyon: Located in Arizona, USA"}, + } + + ids, err := indexer.Store(ctx, docs) + if err != nil { + panic(err) + } + + fmt.Printf("Stored %d documents: %v\n", len(ids), ids) +} + +// mockEmbedding is a placeholder - replace with a real embedding implementation. +type mockEmbedding struct { + dims int +} + +func (m *mockEmbedding) EmbedStrings(_ context.Context, texts []string, _ ...embedding.Option) ([][]float64, error) { + vectors := make([][]float64, len(texts)) + for i := range texts { + vectors[i] = make([]float64, m.dims) + } + return vectors, nil +} diff --git a/components/indexer/valkey/examples/json_indexer/json_indexer.go b/components/indexer/valkey/examples/json_indexer/json_indexer.go new file mode 100644 index 000000000..81bccdc33 --- /dev/null +++ b/components/indexer/valkey/examples/json_indexer/json_indexer.go @@ -0,0 +1,90 @@ +//go:build ignore + +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package main + +import ( + "context" + "fmt" + + glide "github.com/valkey-io/valkey-glide/go/v2" + "github.com/valkey-io/valkey-glide/go/v2/config" + + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/schema" + + vi "github.com/cloudwego/eino-ext/components/indexer/valkey" +) + +// This example demonstrates indexing documents as Valkey JSON. +// Prerequisites: +// - Valkey 9.1+ with Search and JSON modules running on localhost:6379 +// - Create an index first: +// FT.CREATE my_json_index ON JSON PREFIX 1 jdoc: SCHEMA +// $.content AS content TEXT $.vector_content AS vector_content VECTOR HNSW 6 +// TYPE FLOAT32 DIM 1024 DISTANCE_METRIC COSINE +func main() { + ctx := context.Background() + + // 1. Create Valkey GLIDE client + cfg := config.NewClientConfiguration(). + WithAddress(&config.NodeAddress{Host: "localhost", Port: 6379}) + client, err := glide.NewClient(cfg) + if err != nil { + panic(fmt.Sprintf("failed to create client: %v", err)) + } + defer client.Close() + + // 2. Create JSON indexer + indexer, err := vi.NewIndexer(ctx, &vi.IndexerConfig{ + Client: client, + KeyPrefix: "jdoc:", + DocumentType: vi.DocumentTypeJSON, + BatchSize: 10, + Embedding: &mockEmbedding{dims: 1024}, // replace with real embedder + }) + if err != nil { + panic(err) + } + + // 3. Store documents as JSON + docs := []*schema.Document{ + {ID: "1", Content: "Eiffel Tower: Located in Paris, France", MetaData: map[string]any{"category": "europe"}}, + {ID: "2", Content: "Great Wall of China", MetaData: map[string]any{"category": "asia"}}, + } + + ids, err := indexer.Store(ctx, docs) + if err != nil { + panic(err) + } + + fmt.Printf("Stored %d JSON documents: %v\n", len(ids), ids) +} + +// mockEmbedding is a placeholder - replace with a real embedding implementation. +type mockEmbedding struct { + dims int +} + +func (m *mockEmbedding) EmbedStrings(_ context.Context, texts []string, _ ...embedding.Option) ([][]float64, error) { + vectors := make([][]float64, len(texts)) + for i := range texts { + vectors[i] = make([]float64, m.dims) + } + return vectors, nil +} diff --git a/components/indexer/valkey/go.mod b/components/indexer/valkey/go.mod new file mode 100644 index 000000000..d398d35da --- /dev/null +++ b/components/indexer/valkey/go.mod @@ -0,0 +1,40 @@ +module github.com/cloudwego/eino-ext/components/indexer/valkey + +go 1.23.0 + +require ( + github.com/cloudwego/eino v0.6.0 + github.com/valkey-io/valkey-glide/go/v2 v2.4.1 +) + +require ( + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/bytedance/gopkg v0.1.3 // indirect + github.com/bytedance/sonic v1.15.2 // indirect + github.com/bytedance/sonic/loader v0.5.1 // indirect + github.com/cloudwego/base64x v0.1.6 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/eino-contrib/jsonschema v1.0.2 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/goph/emperror v0.17.2 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/cpuid/v2 v2.2.9 // indirect + github.com/kr/text v0.2.0 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/nikolalohinski/gonja v1.5.3 // indirect + github.com/pelletier/go-toml/v2 v2.0.9 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/sirupsen/logrus v1.9.3 // indirect + github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/yargevad/filepathx v1.0.0 // indirect + golang.org/x/arch v0.11.0 // indirect + golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect + golang.org/x/sys v0.26.0 // indirect + google.golang.org/protobuf v1.33.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/components/indexer/valkey/go.sum b/components/indexer/valkey/go.sum new file mode 100644 index 000000000..4432ee7d3 --- /dev/null +++ b/components/indexer/valkey/go.sum @@ -0,0 +1,150 @@ +github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA= +github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/bugsnag/bugsnag-go v1.4.0/go.mod h1:2oa8nejYd4cQ/b0hMIopN0lCRxU0bueqREvZLWFrtK8= +github.com/bugsnag/panicwrap v1.2.0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE= +github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M= +github.com/bytedance/gopkg v0.1.3/go.mod h1:576VvJ+eJgyCzdjS+c4+77QF3p7ubbtiKARP3TxducM= +github.com/bytedance/sonic v1.15.2 h1:90H+rcF/FwLXwfB1cudOLq/je83n683Utf4Cbp0xHCo= +github.com/bytedance/sonic v1.15.2/go.mod h1:mT2NbXunuaEbnZ+mRIX/vYqKISmgEuHFDI4UzmKx2SA= +github.com/bytedance/sonic/loader v0.5.1 h1:Ygpfa9zwRCCKSlrp5bBP/b/Xzc3VxsAW+5NIYXrOOpI= +github.com/bytedance/sonic/loader v0.5.1/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo= +github.com/certifi/gocertifi v0.0.0-20190105021004-abcd57078448/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4= +github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M= +github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU= +github.com/cloudwego/eino v0.6.0 h1:pobGKMOfcQHVNhD9UT/HrvO0eYG6FC2ML/NKY2Eb9+Q= +github.com/cloudwego/eino v0.6.0/go.mod h1:JNapfU+QUrFFpboNDrNOFvmz0m9wjBFHHCr77RH6a50= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +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/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/eino-contrib/jsonschema v1.0.2 h1:HaxruBMUdnXa7Lg/lX8g0Hk71ZIfdTZXmBQz0e3esr8= +github.com/eino-contrib/jsonschema v1.0.2/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4= +github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127 h1:0gkP6mzaMqkmpcJYCFOLkIBwI7xFExG03bbkOkCvUPI= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127/go.mod h1:9ES+weclKsC9YodN5RgxqK/VD9HM9JsCSh7rNhMZE98= +github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/goph/emperror v0.17.2 h1:yLapQcmEsO0ipe9p5TaN22djm3OFV/TfM/fcYP0/J18= +github.com/goph/emperror v0.17.2/go.mod h1:+ZbQ+fUNO/6FNiUo0ujtMjhgad9Xa6fQL9KhH4LNHic= +github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g= +github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo= +github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= +github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8= +github.com/klauspost/cpuid/v2 v2.2.9 h1:66ze0taIn2H33fBvCkXuv9BmCwDfafmiIVpKV9kKGuY= +github.com/klauspost/cpuid/v2 v2.2.9/go.mod h1:rqkxqrZ1EhYM9G+hXH7YdowN5R5RGN6NK4QwQ3WMXF8= +github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mattn/go-colorable v0.1.2 h1:/bC9yWikZXAL9uJdulbSfyVNIR3n3trXl+v8+1sx8mU= +github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE= +github.com/mattn/go-isatty v0.0.8 h1:HLtExJ+uU2HOZ+wI0Tt5DtUDrx8yhUqDcp7fYERX4CE= +github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/nikolalohinski/gonja v1.5.3 h1:GsA+EEaZDZPGJ8JtpeGN78jidhOlxeJROpqMT9fTj9c= +github.com/nikolalohinski/gonja v1.5.3/go.mod h1:RmjwxNiXAEqcq1HeK5SSMmqFJvKOfTfXhkJv6YBtPa4= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/pelletier/go-toml/v2 v2.0.9 h1:uH2qQXheeefCCkuBBSLi7jCiSmj3VRh2+Goq2N7Xxu0= +github.com/pelletier/go-toml/v2 v2.0.9/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc= +github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= +github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= +github.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtISeXco0L5PKQ= +github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f h1:Z2cODYsUxQPofhpYRMQVwWz4yUVpHF+vPi+eUdruUYI= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f/go.mod h1:JqzWyvTuI2X4+9wOHmKSQCYxybB/8j6Ko43qVmXDuZg= +github.com/smarty/assertions v1.15.0 h1:cR//PqUBUiQRakZWqBiFFQ9wb8emQGDb0HeGdqGByCY= +github.com/smarty/assertions v1.15.0/go.mod h1:yABtdzeQs6l1brC900WlRNwj6ZR55d7B+E8C6HtKdec= +github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY= +github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= +github.com/stretchr/testify v1.3.0/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.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/valkey-io/valkey-glide/go/v2 v2.4.1 h1:UUhFmNqeZNSSpM/s4ZeVY6GSFcLdRTyI/ySYEEkxpNY= +github.com/valkey-io/valkey-glide/go/v2 v2.4.1/go.mod h1:LK5zmODJa5xnxZndarh1trntExb3GVGJXz4GwDCagho= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/x-cray/logrus-prefixed-formatter v0.5.2 h1:00txxvfBM9muc0jiLIEAkAcIMJzfthRT6usrui8uGmg= +github.com/x-cray/logrus-prefixed-formatter v0.5.2/go.mod h1:2duySbKsL6M18s5GU7VPsoEPHyzalCE06qoARUCeBBE= +github.com/yargevad/filepathx v1.0.0 h1:SYcT+N3tYGi+NvazubCNlvgIPbzAk7i7y2dwg3I5FYc= +github.com/yargevad/filepathx v1.0.0/go.mod h1:BprfX/gpYNJHJfc35GjRRpVcwWXS89gGulUIU5tK3tA= +go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU= +go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc= +golang.org/x/arch v0.11.0 h1:KXV8WWKCXm6tRpLirl2szsO5j/oOODwZf4hATmGVNs4= +golang.org/x/arch v0.11.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= +golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= +golang.org/x/crypto v0.11.0 h1:6Ewdq3tDic1mg5xRO4milcWCfMVQhI4NkqWWvqejpuA= +golang.org/x/crypto v0.11.0/go.mod h1:xgJhtzW8F9jGdVFWZESrid1U1bjeNy4zgy5cRr/CIio= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 h1:MGwJjxBy0HJshjDNfLsYO8xppfqWlA5ZT9OhtUUhTNw= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo= +golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.10.0 h1:3R7pNqamzBraeqj/Tj8qt1aQ2HpmlC+Cx/qL/7hn4/c= +golang.org/x/term v0.10.0/go.mod h1:lpqdcUyK/oCiQxvxVrppt5ggO2KCZ5QblwqPnfZ6d5o= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI= +google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +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= diff --git a/components/indexer/valkey/indexer.go b/components/indexer/valkey/indexer.go new file mode 100644 index 000000000..7d62461b7 --- /dev/null +++ b/components/indexer/valkey/indexer.go @@ -0,0 +1,408 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package valkey + +import ( + "context" + "encoding/binary" + "encoding/json" + "fmt" + "math" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/schema" +) + +// DocumentType controls how documents are stored in Valkey. +type DocumentType int + +const ( + // DocumentTypeHash stores documents as Valkey Hashes (HSET). + // Index must be created with: FT.CREATE ... ON HASH + DocumentTypeHash DocumentType = iota + // DocumentTypeJSON stores documents as JSON (JSON.SET). + // Requires the Valkey JSON module. + // Index must be created with: FT.CREATE ... ON JSON + DocumentTypeJSON +) + +// BatchClient is the interface for Valkey clients that support pipeline batch execution. +// Each command is a slice of strings (e.g., ["HSET", key, field1, val1, ...]). +// The caller retains ownership and is responsible for closing the client +// after the Indexer is no longer in use. +type BatchClient interface { + Exec(ctx context.Context, commands [][]string) ([]any, error) +} + +// IndexerConfig configures the Valkey indexer. +type IndexerConfig struct { + // Client is a Valkey GLIDE client that supports batch Exec. + // The caller retains ownership and is responsible for closing the client + // after the Indexer is no longer in use. + Client BatchClient + // KeyPrefix is prepended to each document key. + // Ensure this matches the prefix used in FT.CREATE for the search index. + KeyPrefix string + // DocumentType controls storage format. Default: DocumentTypeHash. + DocumentType DocumentType + // DocumentToHashes converts a document into a Hashes struct for Hash storage. + // Only used when DocumentType is DocumentTypeHash. + // Default: stores doc.Content in "content" field with embedding in "vector_content". + DocumentToHashes func(ctx context.Context, doc *schema.Document) (*Hashes, error) + // DocumentToJSON converts a document into a JSON map for JSON storage. + // Only used when DocumentType is DocumentTypeJSON. + // Default: stores content, vector, and metadata fields. + DocumentToJSON func(ctx context.Context, doc *schema.Document, vector []float64) (map[string]any, error) + // BatchSize controls how many texts are embedded per batch call. Default: 10. + BatchSize int + // Embedding is the embedder used to vectorize document content. + Embedding embedding.Embedder +} + +// Hashes represents the key and field-value pairs for a Valkey hash. +type Hashes struct { + // Key is the hash key (without prefix). + Key string + // Field2Value maps field names to their values and embedding configuration. + Field2Value map[string]FieldValue +} + +// FieldValue represents a single field value with optional embedding. +type FieldValue struct { + // Value is the original field value. + Value any + // EmbedKey, if set, indicates Value should be vectorized and stored under this field name. + EmbedKey string + // Stringify converts Value to a string for embedding. If nil, Value is asserted as string. + Stringify func(val any) (string, error) +} + +// Indexer implements indexer.Indexer using Valkey hashes or JSON. +type Indexer struct { + config *IndexerConfig +} + +// NewIndexer creates a new Valkey indexer. +func NewIndexer(_ context.Context, config *IndexerConfig) (*Indexer, error) { + if config.Embedding == nil { + return nil, fmt.Errorf("[NewIndexer] embedding not provided for valkey indexer") + } + if config.Client == nil { + return nil, fmt.Errorf("[NewIndexer] valkey client not provided") + } + cfg := *config // shallow copy to avoid mutating caller's config + if cfg.DocumentToHashes == nil { + cfg.DocumentToHashes = defaultDocumentToFields + } + if cfg.DocumentToJSON == nil { + cfg.DocumentToJSON = defaultDocumentToJSON + } + if cfg.BatchSize == 0 { + cfg.BatchSize = 10 + } + return &Indexer{config: &cfg}, nil +} + +// Store writes documents to Valkey in batches. If a batch fails partway through, +// previously successful batches are NOT rolled back. Callers should treat Store +// as at-least-once and ensure document IDs are deterministic for idempotent writes. +func (i *Indexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) (ids []string, err error) { + options := indexer.GetCommonOptions(&indexer.Options{ + Embedding: i.config.Embedding, + }, opts...) + + ctx = callbacks.EnsureRunInfo(ctx, i.GetType(), components.ComponentOfIndexer) + ctx = callbacks.OnStart(ctx, &indexer.CallbackInput{Docs: docs}) + defer func() { + if err != nil { + callbacks.OnError(ctx, err) + } + }() + + switch i.config.DocumentType { + case DocumentTypeJSON: + err = i.batchJSONSet(ctx, docs, options) + default: + err = i.batchHSet(ctx, docs, options) + } + if err != nil { + return nil, err + } + + ids = make([]string, 0, len(docs)) + for _, doc := range docs { + ids = append(ids, doc.ID) + } + + callbacks.OnEnd(ctx, &indexer.CallbackOutput{IDs: ids}) + return ids, nil +} + +type tuple struct { + key string + fields map[string]string + key2Idx map[string]int +} + +func (i *Indexer) batchHSet(ctx context.Context, docs []*schema.Document, options *indexer.Options) error { + emb := options.Embedding + + var ( + tuples []tuple + texts []string + ) + + embAndStore := func() error { + var vectors [][]float64 + if len(texts) > 0 { + if emb == nil { + return fmt.Errorf("[batchHSet] embedding method not provided") + } + var err error + vectors, err = emb.EmbedStrings(i.makeEmbeddingCtx(ctx, emb), texts) + if err != nil { + return fmt.Errorf("[batchHSet] embedding failed, %w", err) + } + if len(vectors) != len(texts) { + return fmt.Errorf("[batchHSet] invalid vector length, expected=%d, got=%d", len(texts), len(vectors)) + } + } + + commands := make([][]string, 0, len(tuples)) + for _, t := range tuples { + for k, idx := range t.key2Idx { + t.fields[k] = string(vector2Bytes(vectors[idx])) + } + cmd := []string{"HSET", i.config.KeyPrefix + t.key} + for k, v := range t.fields { + cmd = append(cmd, k, v) + } + commands = append(commands, cmd) + } + + if _, err := i.config.Client.Exec(ctx, commands); err != nil { + return err + } + + tuples = tuples[:0] + texts = texts[:0] + return nil + } + + for _, doc := range docs { + hashes, err := i.config.DocumentToHashes(ctx, doc) + if err != nil { + return err + } + + fields := make(map[string]string, len(hashes.Field2Value)) + embSize := 0 + for k, v := range hashes.Field2Value { + fields[k] = fmt.Sprintf("%v", v.Value) + if v.EmbedKey != "" { + embSize++ + } + } + + if embSize > i.config.BatchSize { + return fmt.Errorf("[batchHSet] embedding size over batch size, batch size=%d, got size=%d", + i.config.BatchSize, embSize) + } + + if len(texts)+embSize > i.config.BatchSize { + if err = embAndStore(); err != nil { + return err + } + } + + key2Idx := make(map[string]int, embSize) + for k, v := range hashes.Field2Value { + if v.EmbedKey != "" { + if _, found := fields[v.EmbedKey]; found { + return fmt.Errorf("[batchHSet] duplicate key for value and vector, field=%s", k) + } + var text string + if v.Stringify != nil { + text, err = v.Stringify(v.Value) + if err != nil { + return err + } + } else { + var ok bool + text, ok = v.Value.(string) + if !ok { + return fmt.Errorf("[batchHSet] assert value as string failed, key=%s, emb_key=%s", k, v.EmbedKey) + } + } + key2Idx[v.EmbedKey] = len(texts) + texts = append(texts, text) + } + } + + tuples = append(tuples, tuple{ + key: hashes.Key, + fields: fields, + key2Idx: key2Idx, + }) + } + + if len(tuples) > 0 { + if err := embAndStore(); err != nil { + return err + } + } + + return nil +} + +type jsonDoc struct { + doc *schema.Document + text string +} + +func (i *Indexer) batchJSONSet(ctx context.Context, docs []*schema.Document, options *indexer.Options) error { + emb := options.Embedding + if emb == nil { + return fmt.Errorf("[batchJSONSet] embedding method not provided") + } + + // Collect texts for embedding in batches + var ( + pending []jsonDoc + texts []string + ) + + flushBatch := func() error { + if len(texts) == 0 { + return nil + } + + vectors, err := emb.EmbedStrings(i.makeEmbeddingCtx(ctx, emb), texts) + if err != nil { + return fmt.Errorf("[batchJSONSet] embedding failed, %w", err) + } + if len(vectors) != len(texts) { + return fmt.Errorf("[batchJSONSet] invalid vector length, expected=%d, got=%d", len(texts), len(vectors)) + } + + commands := make([][]string, 0, len(pending)) + for idx, p := range pending { + jsonMap, err := i.config.DocumentToJSON(ctx, p.doc, vectors[idx]) + if err != nil { + return fmt.Errorf("[batchJSONSet] DocumentToJSON failed for key=%s: %w", p.doc.ID, err) + } + + jsonBytes, err := json.Marshal(jsonMap) + if err != nil { + return fmt.Errorf("[batchJSONSet] json marshal failed, %w", err) + } + + commands = append(commands, []string{"JSON.SET", i.config.KeyPrefix + p.doc.ID, "$", string(jsonBytes)}) + } + + if _, err := i.config.Client.Exec(ctx, commands); err != nil { + return err + } + + pending = pending[:0] + texts = texts[:0] + return nil + } + + for _, doc := range docs { + if doc.ID == "" { + return fmt.Errorf("[batchJSONSet] document ID must not be empty") + } + if len(texts) >= i.config.BatchSize { + if err := flushBatch(); err != nil { + return err + } + } + pending = append(pending, jsonDoc{doc: doc, text: doc.Content}) + texts = append(texts, doc.Content) + } + + return flushBatch() +} + +func (i *Indexer) makeEmbeddingCtx(ctx context.Context, emb embedding.Embedder) context.Context { + runInfo := &callbacks.RunInfo{ + Component: components.ComponentOfEmbedding, + } + if embType, ok := components.GetType(emb); ok { + runInfo.Type = embType + } + runInfo.Name = runInfo.Type + string(runInfo.Component) + return callbacks.ReuseHandlers(ctx, runInfo) +} + +const typ = "Valkey" + +func (i *Indexer) GetType() string { + return typ +} + +func (i *Indexer) IsCallbacksEnabled() bool { + return true +} + +func defaultDocumentToFields(_ context.Context, doc *schema.Document) (*Hashes, error) { + if doc.ID == "" { + return nil, fmt.Errorf("[defaultDocumentToFields] doc id not set") + } + field2Value := map[string]FieldValue{ + defaultReturnFieldContent: { + Value: doc.Content, + EmbedKey: defaultReturnFieldVectorContent, + }, + } + for k, v := range doc.MetaData { + if k == defaultReturnFieldContent || k == defaultReturnFieldVectorContent { + return nil, fmt.Errorf("[defaultDocumentToFields] metadata key %q conflicts with reserved field", k) + } + field2Value[k] = FieldValue{Value: v} + } + return &Hashes{Key: doc.ID, Field2Value: field2Value}, nil +} + +func defaultDocumentToJSON(_ context.Context, doc *schema.Document, vector []float64) (map[string]any, error) { + if doc.ID == "" { + return nil, fmt.Errorf("[defaultDocumentToJSON] doc id not set") + } + m := map[string]any{ + defaultReturnFieldContent: doc.Content, + defaultReturnFieldVectorContent: vector, + } + for k, v := range doc.MetaData { + if k == defaultReturnFieldContent || k == defaultReturnFieldVectorContent { + return nil, fmt.Errorf("[defaultDocumentToJSON] metadata key %q conflicts with reserved field", k) + } + m[k] = v + } + return m, nil +} + +func vector2Bytes(vector []float64) []byte { + bytes := make([]byte, len(vector)*4) + for i, v := range vector { + binary.LittleEndian.PutUint32(bytes[i*4:], math.Float32bits(float32(v))) + } + return bytes +} diff --git a/components/indexer/valkey/indexer_integration_test.go b/components/indexer/valkey/indexer_integration_test.go new file mode 100644 index 000000000..c7db52211 --- /dev/null +++ b/components/indexer/valkey/indexer_integration_test.go @@ -0,0 +1,174 @@ +//go:build integration + +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package valkey + +import ( + "context" + "net" + "os" + "strconv" + "testing" + + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/schema" + glide "github.com/valkey-io/valkey-glide/go/v2" + "github.com/valkey-io/valkey-glide/go/v2/config" +) + +func getTestClient(t *testing.T) *glide.Client { + t.Helper() + addr := os.Getenv("VALKEY_ADDR") + if addr == "" { + addr = "localhost:6379" + } + host, portStr, err := net.SplitHostPort(addr) + if err != nil { + t.Fatalf("invalid VALKEY_ADDR %q: %v", addr, err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + t.Fatalf("invalid port in VALKEY_ADDR %q: %v", addr, err) + } + cfg := config.NewClientConfiguration(). + WithAddress(&config.NodeAddress{Host: host, Port: port}) + client, err := glide.NewClient(cfg) + if err != nil { + t.Fatalf("failed to create valkey client: %v", err) + } + return client +} + +type deterministicEmbedder struct { + dim int +} + +func (e *deterministicEmbedder) EmbedStrings(_ context.Context, texts []string, _ ...embedding.Option) ([][]float64, error) { + vectors := make([][]float64, len(texts)) + for i, text := range texts { + vec := make([]float64, e.dim) + for j := 0; j < e.dim; j++ { + if j < len(text) { + vec[j] = float64(text[j]) / 255.0 + } + } + vectors[i] = vec + } + return vectors, nil +} + +func TestIntegration_Indexer_Hash(t *testing.T) { + ctx := context.Background() + client := getTestClient(t) + defer client.Close() + + prefix := "inttest:indexer:hash:" + + t.Cleanup(func() { + client.CustomCommand(ctx, []string{"DEL", prefix + "doc1", prefix + "doc2", prefix + "doc3"}) + }) + + idx, err := NewIndexer(ctx, &IndexerConfig{ + Client: client, + KeyPrefix: prefix, + DocumentType: DocumentTypeHash, + BatchSize: 10, + Embedding: &deterministicEmbedder{dim: 8}, + }) + if err != nil { + t.Fatalf("NewIndexer failed: %v", err) + } + + docs := []*schema.Document{ + {ID: "doc1", Content: "Valkey is a high-performance key-value store"}, + {ID: "doc2", Content: "Vector search enables semantic similarity queries"}, + {ID: "doc3", Content: "The quick brown fox", MetaData: map[string]any{"category": "example"}}, + } + + ids, err := idx.Store(ctx, docs) + if err != nil { + t.Fatalf("Store failed: %v", err) + } + if len(ids) != 3 { + t.Fatalf("expected 3 ids, got %d", len(ids)) + } + + // Verify via HGetAll + for _, doc := range docs { + result, err := client.HGetAll(ctx, prefix+doc.ID) + if err != nil { + t.Fatalf("HGetAll failed for %s: %v", doc.ID, err) + } + if result[defaultReturnFieldContent] != doc.Content { + t.Fatalf("content mismatch for %s", doc.ID) + } + if result[defaultReturnFieldVectorContent] == "" { + t.Fatalf("vector not set for %s", doc.ID) + } + } + t.Logf("Successfully stored %d hash documents", len(ids)) +} + +func TestIntegration_Indexer_JSON(t *testing.T) { + ctx := context.Background() + client := getTestClient(t) + defer client.Close() + + prefix := "inttest:indexer:json:" + + t.Cleanup(func() { + client.CustomCommand(ctx, []string{"DEL", prefix + "doc1", prefix + "doc2"}) + }) + + idx, err := NewIndexer(ctx, &IndexerConfig{ + Client: client, + KeyPrefix: prefix, + DocumentType: DocumentTypeJSON, + BatchSize: 10, + Embedding: &deterministicEmbedder{dim: 8}, + }) + if err != nil { + t.Fatalf("NewIndexer failed: %v", err) + } + + docs := []*schema.Document{ + {ID: "doc1", Content: "JSON document storage in Valkey"}, + {ID: "doc2", Content: "Supports nested structures", MetaData: map[string]any{"tag": "test"}}, + } + + ids, err := idx.Store(ctx, docs) + if err != nil { + t.Fatalf("Store failed: %v", err) + } + if len(ids) != 2 { + t.Fatalf("expected 2 ids, got %d", len(ids)) + } + + // Verify via JSON.GET + for _, doc := range docs { + result, err := client.CustomCommand(ctx, []string{"JSON.GET", prefix + doc.ID, "$." + defaultReturnFieldContent}) + if err != nil { + t.Fatalf("JSON.GET failed for %s: %v", doc.ID, err) + } + if result == nil { + t.Fatalf("no JSON data for %s", doc.ID) + } + t.Logf(" %s: %v", doc.ID, result) + } + t.Logf("Successfully stored %d JSON documents", len(ids)) +} diff --git a/components/indexer/valkey/indexer_test.go b/components/indexer/valkey/indexer_test.go new file mode 100644 index 000000000..fd8e4a7de --- /dev/null +++ b/components/indexer/valkey/indexer_test.go @@ -0,0 +1,397 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package valkey + +import ( + "context" + "fmt" + "testing" + + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/schema" +) + +type mockBatchClient struct { + execCount int + execErr error +} + +func (m *mockBatchClient) Exec(_ context.Context, _ [][]string) ([]any, error) { + m.execCount++ + if m.execErr != nil { + return nil, m.execErr + } + return []any{}, nil +} + +type mockEmbedding struct { + err error + cnt int + sizeForCall []int + dims int +} + +func (m *mockEmbedding) EmbedStrings(_ context.Context, texts []string, _ ...embedding.Option) ([][]float64, error) { + if m.err != nil { + return nil, m.err + } + if m.cnt >= len(m.sizeForCall) { + return nil, fmt.Errorf("unexpected call") + } + slice := make([]float64, m.dims) + for i := range slice { + slice[i] = 1.1 + } + r := make([][]float64, m.sizeForCall[m.cnt]) + m.cnt++ + for i := range r { + r[i] = slice + } + return r, nil +} + +func TestNewIndexer(t *testing.T) { + ctx := context.Background() + client := &mockBatchClient{} + + tests := []struct { + name string + config *IndexerConfig + wantErr string + }{ + { + name: "embedding not provided", + config: &IndexerConfig{Client: client}, + wantErr: "[NewIndexer] embedding not provided for valkey indexer", + }, + { + name: "client not provided", + config: &IndexerConfig{Embedding: &mockEmbedding{}}, + wantErr: "[NewIndexer] valkey client not provided", + }, + { + name: "success", + config: &IndexerConfig{Client: client, Embedding: &mockEmbedding{}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + idx, err := NewIndexer(ctx, tt.config) + if tt.wantErr != "" { + if err == nil || err.Error() != tt.wantErr { + t.Fatalf("expected error %q, got %v", tt.wantErr, err) + } + if idx != nil { + t.Fatal("expected nil indexer on error") + } + } else { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if idx == nil { + t.Fatal("expected non-nil indexer") + } + } + }) + } +} + +func TestNewIndexer_DoesNotMutateConfig(t *testing.T) { + ctx := context.Background() + cfg := &IndexerConfig{ + Client: &mockBatchClient{}, + Embedding: &mockEmbedding{}, + } + _, err := NewIndexer(ctx, cfg) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.BatchSize != 0 { + t.Fatal("NewIndexer mutated caller's config") + } +} + +func TestIndexer_Store_Hash_Success(t *testing.T) { + ctx := context.Background() + client := &mockBatchClient{} + + idx := &Indexer{config: &IndexerConfig{ + Client: client, + KeyPrefix: "test:", + DocumentType: DocumentTypeHash, + DocumentToHashes: defaultDocumentToFields, + BatchSize: 10, + Embedding: &mockEmbedding{sizeForCall: []int{2}, dims: 4}, + }} + + docs := []*schema.Document{ + {ID: "1", Content: "hello"}, + {ID: "2", Content: "world", MetaData: map[string]any{"tag": "test"}}, + } + + ids, err := idx.Store(ctx, docs) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(ids) != 2 || ids[0] != "1" || ids[1] != "2" { + t.Fatalf("unexpected ids: %v", ids) + } + if client.execCount != 1 { + t.Fatalf("expected 1 batch exec call, got %d", client.execCount) + } +} + +func TestIndexer_Store_JSON_Success(t *testing.T) { + ctx := context.Background() + client := &mockBatchClient{} + + idx := &Indexer{config: &IndexerConfig{ + Client: client, + KeyPrefix: "json:", + DocumentType: DocumentTypeJSON, + DocumentToJSON: defaultDocumentToJSON, + BatchSize: 10, + Embedding: &mockEmbedding{sizeForCall: []int{2}, dims: 4}, + }} + + docs := []*schema.Document{ + {ID: "1", Content: "hello"}, + {ID: "2", Content: "world", MetaData: map[string]any{"tag": "test"}}, + } + + ids, err := idx.Store(ctx, docs) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(ids) != 2 || ids[0] != "1" || ids[1] != "2" { + t.Fatalf("unexpected ids: %v", ids) + } + if client.execCount != 1 { + t.Fatalf("expected 1 batch exec call, got %d", client.execCount) + } +} + +func TestIndexer_Store_EmbeddingError(t *testing.T) { + ctx := context.Background() + client := &mockBatchClient{} + + idx := &Indexer{config: &IndexerConfig{ + Client: client, + DocumentType: DocumentTypeHash, + DocumentToHashes: defaultDocumentToFields, + BatchSize: 10, + Embedding: &mockEmbedding{err: fmt.Errorf("embed error")}, + }} + + docs := []*schema.Document{{ID: "1", Content: "hello"}} + _, err := idx.Store(ctx, docs) + if err == nil { + t.Fatal("expected error") + } +} + +func TestIndexer_Store_BatchSizeExceeded(t *testing.T) { + ctx := context.Background() + client := &mockBatchClient{} + + idx := &Indexer{config: &IndexerConfig{ + Client: client, + DocumentToHashes: func(_ context.Context, doc *schema.Document) (*Hashes, error) { + return &Hashes{ + Key: doc.ID, + Field2Value: map[string]FieldValue{ + "f1": {Value: "v1", EmbedKey: "e1"}, + "f2": {Value: "v2", EmbedKey: "e2"}, + }, + }, nil + }, + BatchSize: 1, + Embedding: &mockEmbedding{}, + }} + + docs := []*schema.Document{{ID: "1", Content: "hello"}} + _, err := idx.Store(ctx, docs) + if err == nil { + t.Fatal("expected batch size error") + } +} + +func TestIndexer_Store_DocumentToHashesError(t *testing.T) { + ctx := context.Background() + client := &mockBatchClient{} + + idx := &Indexer{config: &IndexerConfig{ + Client: client, + DocumentToHashes: func(_ context.Context, _ *schema.Document) (*Hashes, error) { + return nil, fmt.Errorf("conversion error") + }, + BatchSize: 10, + Embedding: &mockEmbedding{}, + }} + + docs := []*schema.Document{{ID: "1", Content: "hello"}} + _, err := idx.Store(ctx, docs) + if err == nil || err.Error() != "conversion error" { + t.Fatalf("expected conversion error, got: %v", err) + } +} + +func TestDefaultDocumentToFields(t *testing.T) { + t.Run("success", func(t *testing.T) { + doc := &schema.Document{ID: "test", Content: "hello", MetaData: map[string]any{"k": "v"}} + h, err := defaultDocumentToFields(context.Background(), doc) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if h.Key != "test" { + t.Fatalf("unexpected key: %s", h.Key) + } + if h.Field2Value[defaultReturnFieldContent].EmbedKey != defaultReturnFieldVectorContent { + t.Fatal("unexpected embed key") + } + }) + + t.Run("empty id", func(t *testing.T) { + doc := &schema.Document{Content: "hello"} + _, err := defaultDocumentToFields(context.Background(), doc) + if err == nil { + t.Fatal("expected error for empty id") + } + }) + + t.Run("metadata conflicts with reserved field", func(t *testing.T) { + doc := &schema.Document{ID: "test", Content: "hello", MetaData: map[string]any{"content": "override"}} + _, err := defaultDocumentToFields(context.Background(), doc) + if err == nil { + t.Fatal("expected error for reserved field conflict") + } + }) +} + +func TestDefaultDocumentToJSON(t *testing.T) { + doc := &schema.Document{ID: "test", Content: "hello", MetaData: map[string]any{"k": "v"}} + vec := []float64{1.0, 2.0} + m, err := defaultDocumentToJSON(context.Background(), doc, vec) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if m[defaultReturnFieldContent] != "hello" { + t.Fatal("content mismatch") + } + if m["k"] != "v" { + t.Fatal("metadata mismatch") + } + vecResult := m[defaultReturnFieldVectorContent].([]float64) + if len(vecResult) != 2 { + t.Fatal("vector mismatch") + } +} + +func TestDefaultDocumentToJSON_ReservedFieldConflict(t *testing.T) { + doc := &schema.Document{ID: "test", Content: "hello", MetaData: map[string]any{"vector_content": "bad"}} + _, err := defaultDocumentToJSON(context.Background(), doc, []float64{1.0}) + if err == nil { + t.Fatal("expected error for reserved field conflict") + } +} + +func TestIndexer_GetType(t *testing.T) { + idx := &Indexer{} + if idx.GetType() != "Valkey" { + t.Fatalf("expected Valkey, got %s", idx.GetType()) + } +} + +func TestIndexer_IsCallbacksEnabled(t *testing.T) { + idx := &Indexer{} + if !idx.IsCallbacksEnabled() { + t.Fatal("expected callbacks enabled") + } +} + +func TestIndexer_Store_WithOptions(t *testing.T) { + ctx := context.Background() + client := &mockBatchClient{} + + idx := &Indexer{config: &IndexerConfig{ + Client: client, + DocumentToHashes: defaultDocumentToFields, + BatchSize: 10, + Embedding: &mockEmbedding{sizeForCall: []int{1}, dims: 4}, + }} + + docs := []*schema.Document{{ID: "1", Content: "hello"}} + ids, err := idx.Store(ctx, docs, indexer.WithEmbedding(&mockEmbedding{sizeForCall: []int{1}, dims: 4})) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(ids) != 1 || ids[0] != "1" { + t.Fatalf("unexpected ids: %v", ids) + } +} + +func TestIndexer_Store_MultiBatch(t *testing.T) { + ctx := context.Background() + client := &mockBatchClient{} + + idx := &Indexer{config: &IndexerConfig{ + Client: client, + DocumentToHashes: defaultDocumentToFields, + BatchSize: 2, // Force multiple batches + Embedding: &mockEmbedding{sizeForCall: []int{2, 2, 1}, dims: 4}, + }} + + docs := []*schema.Document{ + {ID: "1", Content: "a"}, + {ID: "2", Content: "b"}, + {ID: "3", Content: "c"}, + {ID: "4", Content: "d"}, + {ID: "5", Content: "e"}, + } + + ids, err := idx.Store(ctx, docs) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(ids) != 5 { + t.Fatalf("expected 5 ids, got %d", len(ids)) + } + if client.execCount != 3 { + t.Fatalf("expected 3 batch exec calls, got %d", client.execCount) + } +} + +func TestIndexer_Store_JSON_EmptyID(t *testing.T) { + ctx := context.Background() + client := &mockBatchClient{} + + idx := &Indexer{config: &IndexerConfig{ + Client: client, + DocumentType: DocumentTypeJSON, + DocumentToJSON: defaultDocumentToJSON, + BatchSize: 10, + Embedding: &mockEmbedding{sizeForCall: []int{1}, dims: 4}, + }} + + docs := []*schema.Document{{ID: "", Content: "hello"}} + _, err := idx.Store(ctx, docs) + if err == nil { + t.Fatal("expected error for empty doc ID") + } +} diff --git a/components/retriever/valkey/README.md b/components/retriever/valkey/README.md new file mode 100644 index 000000000..70b306f75 --- /dev/null +++ b/components/retriever/valkey/README.md @@ -0,0 +1,147 @@ +# Valkey Retriever + +A Valkey retriever implementation for [Eino](https://github.com/cloudwego/eino) that implements the `Retriever` interface. This component uses Valkey Search capabilities (FT.SEARCH) via the [valkey-glide](https://github.com/valkey-io/valkey-glide) client to retrieve documents based on semantic similarity. + +## Features + +- Implements `github.com/cloudwego/eino/components/retriever.Retriever` +- Two search modes: + - KNN vector search for top-k results + - Vector range search with distance threshold +- Hybrid search with filter expressions +- Configurable distance metric, vector field, return fields, dialect +- Embedding integration for automatic query vectorization +- Full Eino callback integration (OnStart, OnEnd, OnError) + +## Installation + +```bash +go get github.com/cloudwego/eino-ext/components/retriever/valkey@latest +``` + +## Prerequisites + +- Valkey 9.1+ with the Search module (e.g., `valkey/valkey-bundle`) +- [valkey-glide](https://github.com/valkey-io/valkey-glide) Go client (requires CGO and Rust core library) + +## Quick Start + +```go +import ( + "context" + "fmt" + + glide "github.com/valkey-io/valkey-glide/go/v2" + "github.com/valkey-io/valkey-glide/go/v2/config" + + valkeyRetriever "github.com/cloudwego/eino-ext/components/retriever/valkey" +) + +func main() { + ctx := context.Background() + + // 1. Create Valkey GLIDE client + cfg := config.NewClientConfiguration(). + WithAddress(&config.NodeAddress{Host: "localhost", Port: 6379}) + client, _ := glide.NewClient(cfg) + + // 2. Create embedding component (use your preferred embedder) + emb := yourEmbedder() + + // 3. Create Valkey retriever with KNN search + retriever, _ := valkeyRetriever.NewRetriever(ctx, &valkeyRetriever.RetrieverConfig{ + Client: client, + Index: "my_index", + VectorField: "vector_content", + TopK: 5, + Embedding: emb, + }) + + // 4. Retrieve documents + docs, err := retriever.Retrieve(ctx, "search query") + if err != nil { + fmt.Printf("retrieve error: %v\n", err) + return + } + + for _, doc := range docs { + fmt.Printf("ID: %s, Content: %s\n", doc.ID, doc.Content) + } +} +``` + +## Configuration + +```go +type RetrieverConfig struct { + // Required: Valkey GLIDE client (must implement CustomCommand) + Client SearchClient + + // Required: Index name for vector search + Index string + + // Optional: Vector field name (default: "vector_content") + VectorField string + + // Optional: Distance threshold for range search + // If set: uses vector range search (requires Valkey Search 1.3/2.0+, not yet supported) + // If nil: uses KNN vector search (default) + DistanceThreshold *float64 + + // Optional: Query dialect (default: 2) + Dialect int + + // Optional: Fields to return (default: ["content", "vector_content"]) + ReturnFields []string + + // Optional: Custom document converter + DocumentConverter func(ctx context.Context, doc models.FtSearchDocument) (*schema.Document, error) + + // Optional: Number of results (default: 5) + TopK int + + // Required: Embedding method for query vectorization + Embedding embedding.Embedder +} +``` + +## Search Modes + +### KNN Vector Search + +```go +retriever, _ := valkeyRetriever.NewRetriever(ctx, &valkeyRetriever.RetrieverConfig{ + Client: client, + Index: "my_index", + TopK: 10, + Embedding: emb, +}) +``` + +### Vector Range Search + +> **Note:** VECTOR_RANGE requires Valkey Search 1.3/2.0+, which is not yet released. +> Using this option with current Valkey Search versions will result in an error. + +```go +threshold := 0.5 +retriever, _ := valkeyRetriever.NewRetriever(ctx, &valkeyRetriever.RetrieverConfig{ + Client: client, + Index: "my_index", + DistanceThreshold: &threshold, + Embedding: emb, +}) +``` + +## With Filters + +```go +docs, _ := retriever.Retrieve(ctx, "search query", + valkeyRetriever.WithFilterQuery("@category:{technology}")) +``` + +## For More Details + +- [Eino Documentation](https://www.cloudwego.io/zh/docs/eino/) +- [Valkey GLIDE Go Client](https://github.com/valkey-io/valkey-glide) +- [Valkey Search](https://valkey.io/topics/search/) diff --git a/components/retriever/valkey/consts.go b/components/retriever/valkey/consts.go new file mode 100644 index 000000000..eee41bd58 --- /dev/null +++ b/components/retriever/valkey/consts.go @@ -0,0 +1,30 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package valkey + +const ( + defaultReturnFieldContent = "content" + defaultReturnFieldVectorContent = "vector_content" + paramVector = "vector" + paramDistanceThreshold = "distance_threshold" + // SortByDistanceAttributeName is the alias assigned to the computed distance + // in KNN vector search queries (via "AS distance" in the query expression). + // If a document field has this same name, the distance value will shadow it in + // search results, causing incorrect data in the returned documents. To avoid + // collisions, rename the document field or use a custom DocumentConverter. + SortByDistanceAttributeName = "distance" +) diff --git a/components/retriever/valkey/examples/knn_retriever/knn_retriever.go b/components/retriever/valkey/examples/knn_retriever/knn_retriever.go new file mode 100644 index 000000000..bb28671e5 --- /dev/null +++ b/components/retriever/valkey/examples/knn_retriever/knn_retriever.go @@ -0,0 +1,96 @@ +//go:build ignore + +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package main + +import ( + "context" + "fmt" + + glide "github.com/valkey-io/valkey-glide/go/v2" + "github.com/valkey-io/valkey-glide/go/v2/config" + + "github.com/cloudwego/eino/components/embedding" + + vr "github.com/cloudwego/eino-ext/components/retriever/valkey" +) + +// This example demonstrates using the Valkey retriever for KNN vector search. +// Prerequisites: +// - Valkey 9.1+ with Search module running on localhost:6379 +// - An index created with: FT.CREATE my_index ON HASH PREFIX 1 doc: SCHEMA +// content TEXT vector_content VECTOR HNSW 6 TYPE FLOAT32 DIM 1024 DISTANCE_METRIC COSINE +// - Documents indexed (see indexer example) +func main() { + ctx := context.Background() + + // 1. Create Valkey GLIDE client + cfg := config.NewClientConfiguration(). + WithAddress(&config.NodeAddress{Host: "localhost", Port: 6379}) + client, err := glide.NewClient(cfg) + if err != nil { + panic(fmt.Sprintf("failed to create client: %v", err)) + } + defer client.Close() + + // 2. Create retriever with your embedding model + r, err := vr.NewRetriever(ctx, &vr.RetrieverConfig{ + Client: client, + Index: "my_index", + TopK: 5, + Embedding: &mockEmbedding{dims: 1024}, // replace with real embedder + }) + if err != nil { + panic(err) + } + + // 3. Retrieve documents + docs, err := r.Retrieve(ctx, "tourist attractions in Europe") + if err != nil { + panic(err) + } + + for _, doc := range docs { + fmt.Printf("ID: %s, Content: %s\n", doc.ID, doc.Content) + } + + // 4. Retrieve with filter + docs, err = r.Retrieve(ctx, "tourist attractions", + vr.WithFilterQuery("@category:{europe}")) + if err != nil { + panic(err) + } + + fmt.Printf("\nFiltered results: %d\n", len(docs)) + for _, doc := range docs { + fmt.Printf("ID: %s, Content: %s\n", doc.ID, doc.Content) + } +} + +// mockEmbedding is a placeholder - replace with a real embedding implementation. +type mockEmbedding struct { + dims int +} + +func (m *mockEmbedding) EmbedStrings(_ context.Context, texts []string, _ ...embedding.Option) ([][]float64, error) { + vectors := make([][]float64, len(texts)) + for i := range texts { + vectors[i] = make([]float64, m.dims) + } + return vectors, nil +} diff --git a/components/retriever/valkey/go.mod b/components/retriever/valkey/go.mod new file mode 100644 index 000000000..13487f982 --- /dev/null +++ b/components/retriever/valkey/go.mod @@ -0,0 +1,39 @@ +module github.com/cloudwego/eino-ext/components/retriever/valkey + +go 1.23.0 + +require ( + github.com/cloudwego/eino v0.6.0 + github.com/valkey-io/valkey-glide/go/v2 v2.4.1 +) + +require ( + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/bytedance/gopkg v0.1.3 // indirect + github.com/bytedance/sonic v1.15.2 // indirect + github.com/bytedance/sonic/loader v0.5.1 // indirect + github.com/cloudwego/base64x v0.1.6 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/eino-contrib/jsonschema v1.0.2 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/goph/emperror v0.17.2 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/cpuid/v2 v2.2.9 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/nikolalohinski/gonja v1.5.3 // indirect + github.com/pelletier/go-toml/v2 v2.0.9 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/sirupsen/logrus v1.9.3 // indirect + github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/yargevad/filepathx v1.0.0 // indirect + golang.org/x/arch v0.11.0 // indirect + golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect + golang.org/x/sys v0.26.0 // indirect + google.golang.org/protobuf v1.33.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/components/retriever/valkey/go.sum b/components/retriever/valkey/go.sum new file mode 100644 index 000000000..09bf13657 --- /dev/null +++ b/components/retriever/valkey/go.sum @@ -0,0 +1,148 @@ +github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA= +github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/bugsnag/bugsnag-go v1.4.0/go.mod h1:2oa8nejYd4cQ/b0hMIopN0lCRxU0bueqREvZLWFrtK8= +github.com/bugsnag/panicwrap v1.2.0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE= +github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M= +github.com/bytedance/gopkg v0.1.3/go.mod h1:576VvJ+eJgyCzdjS+c4+77QF3p7ubbtiKARP3TxducM= +github.com/bytedance/sonic v1.15.2 h1:90H+rcF/FwLXwfB1cudOLq/je83n683Utf4Cbp0xHCo= +github.com/bytedance/sonic v1.15.2/go.mod h1:mT2NbXunuaEbnZ+mRIX/vYqKISmgEuHFDI4UzmKx2SA= +github.com/bytedance/sonic/loader v0.5.1 h1:Ygpfa9zwRCCKSlrp5bBP/b/Xzc3VxsAW+5NIYXrOOpI= +github.com/bytedance/sonic/loader v0.5.1/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo= +github.com/certifi/gocertifi v0.0.0-20190105021004-abcd57078448/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4= +github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M= +github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU= +github.com/cloudwego/eino v0.6.0 h1:pobGKMOfcQHVNhD9UT/HrvO0eYG6FC2ML/NKY2Eb9+Q= +github.com/cloudwego/eino v0.6.0/go.mod h1:JNapfU+QUrFFpboNDrNOFvmz0m9wjBFHHCr77RH6a50= +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/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/eino-contrib/jsonschema v1.0.2 h1:HaxruBMUdnXa7Lg/lX8g0Hk71ZIfdTZXmBQz0e3esr8= +github.com/eino-contrib/jsonschema v1.0.2/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4= +github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127 h1:0gkP6mzaMqkmpcJYCFOLkIBwI7xFExG03bbkOkCvUPI= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127/go.mod h1:9ES+weclKsC9YodN5RgxqK/VD9HM9JsCSh7rNhMZE98= +github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/goph/emperror v0.17.2 h1:yLapQcmEsO0ipe9p5TaN22djm3OFV/TfM/fcYP0/J18= +github.com/goph/emperror v0.17.2/go.mod h1:+ZbQ+fUNO/6FNiUo0ujtMjhgad9Xa6fQL9KhH4LNHic= +github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g= +github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo= +github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= +github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8= +github.com/klauspost/cpuid/v2 v2.2.9 h1:66ze0taIn2H33fBvCkXuv9BmCwDfafmiIVpKV9kKGuY= +github.com/klauspost/cpuid/v2 v2.2.9/go.mod h1:rqkxqrZ1EhYM9G+hXH7YdowN5R5RGN6NK4QwQ3WMXF8= +github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mattn/go-colorable v0.1.2 h1:/bC9yWikZXAL9uJdulbSfyVNIR3n3trXl+v8+1sx8mU= +github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE= +github.com/mattn/go-isatty v0.0.8 h1:HLtExJ+uU2HOZ+wI0Tt5DtUDrx8yhUqDcp7fYERX4CE= +github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/nikolalohinski/gonja v1.5.3 h1:GsA+EEaZDZPGJ8JtpeGN78jidhOlxeJROpqMT9fTj9c= +github.com/nikolalohinski/gonja v1.5.3/go.mod h1:RmjwxNiXAEqcq1HeK5SSMmqFJvKOfTfXhkJv6YBtPa4= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/pelletier/go-toml/v2 v2.0.9 h1:uH2qQXheeefCCkuBBSLi7jCiSmj3VRh2+Goq2N7Xxu0= +github.com/pelletier/go-toml/v2 v2.0.9/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc= +github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= +github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= +github.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtISeXco0L5PKQ= +github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f h1:Z2cODYsUxQPofhpYRMQVwWz4yUVpHF+vPi+eUdruUYI= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f/go.mod h1:JqzWyvTuI2X4+9wOHmKSQCYxybB/8j6Ko43qVmXDuZg= +github.com/smarty/assertions v1.15.0 h1:cR//PqUBUiQRakZWqBiFFQ9wb8emQGDb0HeGdqGByCY= +github.com/smarty/assertions v1.15.0/go.mod h1:yABtdzeQs6l1brC900WlRNwj6ZR55d7B+E8C6HtKdec= +github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY= +github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= +github.com/stretchr/testify v1.3.0/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.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/valkey-io/valkey-glide/go/v2 v2.4.1 h1:UUhFmNqeZNSSpM/s4ZeVY6GSFcLdRTyI/ySYEEkxpNY= +github.com/valkey-io/valkey-glide/go/v2 v2.4.1/go.mod h1:LK5zmODJa5xnxZndarh1trntExb3GVGJXz4GwDCagho= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/x-cray/logrus-prefixed-formatter v0.5.2 h1:00txxvfBM9muc0jiLIEAkAcIMJzfthRT6usrui8uGmg= +github.com/x-cray/logrus-prefixed-formatter v0.5.2/go.mod h1:2duySbKsL6M18s5GU7VPsoEPHyzalCE06qoARUCeBBE= +github.com/yargevad/filepathx v1.0.0 h1:SYcT+N3tYGi+NvazubCNlvgIPbzAk7i7y2dwg3I5FYc= +github.com/yargevad/filepathx v1.0.0/go.mod h1:BprfX/gpYNJHJfc35GjRRpVcwWXS89gGulUIU5tK3tA= +go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU= +go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc= +golang.org/x/arch v0.11.0 h1:KXV8WWKCXm6tRpLirl2szsO5j/oOODwZf4hATmGVNs4= +golang.org/x/arch v0.11.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= +golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= +golang.org/x/crypto v0.11.0 h1:6Ewdq3tDic1mg5xRO4milcWCfMVQhI4NkqWWvqejpuA= +golang.org/x/crypto v0.11.0/go.mod h1:xgJhtzW8F9jGdVFWZESrid1U1bjeNy4zgy5cRr/CIio= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 h1:MGwJjxBy0HJshjDNfLsYO8xppfqWlA5ZT9OhtUUhTNw= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo= +golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.10.0 h1:3R7pNqamzBraeqj/Tj8qt1aQ2HpmlC+Cx/qL/7hn4/c= +golang.org/x/term v0.10.0/go.mod h1:lpqdcUyK/oCiQxvxVrppt5ggO2KCZ5QblwqPnfZ6d5o= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI= +google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +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= diff --git a/components/retriever/valkey/options.go b/components/retriever/valkey/options.go new file mode 100644 index 000000000..1c0bbae87 --- /dev/null +++ b/components/retriever/valkey/options.go @@ -0,0 +1,33 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package valkey + +import ( + "github.com/cloudwego/eino/components/retriever" +) + +type implOptions struct { + FilterQuery string +} + +// WithFilterQuery sets a Valkey Search filter expression for hybrid search. +// see: https://valkey.io/docs/interact/search-and-query/advanced-concepts/vectors/#filters +func WithFilterQuery(filter string) retriever.Option { + return retriever.WrapImplSpecificOptFn(func(o *implOptions) { + o.FilterQuery = filter + }) +} diff --git a/components/retriever/valkey/retriever.go b/components/retriever/valkey/retriever.go new file mode 100644 index 000000000..35b972d0c --- /dev/null +++ b/components/retriever/valkey/retriever.go @@ -0,0 +1,329 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package valkey + +import ( + "context" + "encoding/binary" + "fmt" + "math" + "strconv" + "strings" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" +) + +// FtSearchDocument represents a single document returned from FT.SEARCH. +type FtSearchDocument struct { + Key string + Fields map[string]any +} + +// SearchClient is the interface for Valkey clients that support FT.SEARCH. +// The caller retains ownership of the client and is responsible for closing it +// after the Retriever is no longer in use. +// Satisfied by *glide.Client. +type SearchClient interface { + CustomCommand(ctx context.Context, args []string) (any, error) +} + +// RetrieverConfig configures the Valkey retriever. +type RetrieverConfig struct { + // Client is a Valkey GLIDE client that supports FT.SEARCH via CustomCommand. + // The caller retains ownership and is responsible for closing the client + // after the Retriever is no longer in use. + Client SearchClient + // Index is the name of the Valkey Search index to query. + Index string + // VectorField is the vector field name in the index schema. + // Default: "vector_content" + VectorField string + // DistanceThreshold enables vector range queries when set. + // If nil, KNN vector search is used instead. + // NOTE: VECTOR_RANGE requires Valkey Search 1.3/2.0+ which is not yet released. + // Using this option with current Valkey Search versions will result in an error. + DistanceThreshold *float64 + // Dialect is the query dialect version. Default: 2. + Dialect int + // ReturnFields limits which fields are returned from search results. + // Default: []string{"content", "vector_content"} + ReturnFields []string + // DocumentConverter converts a search result document to an eino Document. + // Default: built-in parser using ReturnFields. + DocumentConverter func(ctx context.Context, doc FtSearchDocument) (*schema.Document, error) + // TopK limits the number of results. Default: 5. + TopK int + // Embedding is the embedder used to vectorize the query string. + Embedding embedding.Embedder +} + +// Retriever implements retriever.Retriever using Valkey Search. +type Retriever struct { + config *RetrieverConfig +} + +// NewRetriever creates a new Valkey retriever. +func NewRetriever(_ context.Context, config *RetrieverConfig) (*Retriever, error) { + if config.Embedding == nil { + return nil, fmt.Errorf("[NewRetriever] embedding not provided for valkey retriever") + } + if config.Index == "" { + return nil, fmt.Errorf("[NewRetriever] valkey index not provided") + } + if config.Client == nil { + return nil, fmt.Errorf("[NewRetriever] valkey client not provided") + } + cfg := *config // shallow copy to avoid mutating caller's config + if cfg.Dialect < 2 { + cfg.Dialect = 2 + } + if cfg.TopK == 0 { + cfg.TopK = 5 + } + if cfg.VectorField == "" { + cfg.VectorField = defaultReturnFieldVectorContent + } + if len(cfg.ReturnFields) == 0 { + cfg.ReturnFields = []string{ + defaultReturnFieldContent, + defaultReturnFieldVectorContent, + } + } + if cfg.DocumentConverter == nil { + cfg.DocumentConverter = defaultResultParser(cfg.ReturnFields) + } + return &Retriever{config: &cfg}, nil +} + +// validateFilterQuery rejects filter queries that contain dangerous FT.SEARCH +// syntax which could be used for injection attacks. +func validateFilterQuery(filter string) error { + if strings.Contains(filter, "=>") || strings.Contains(filter, "[KNN") { + return fmt.Errorf("[valkey retriever] filter contains disallowed syntax") + } + return nil +} + +func (r *Retriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) (docs []*schema.Document, err error) { + co := retriever.GetCommonOptions(&retriever.Options{ + Index: &r.config.Index, + TopK: &r.config.TopK, + ScoreThreshold: r.config.DistanceThreshold, + Embedding: r.config.Embedding, + }, opts...) + io := retriever.GetImplSpecificOptions(&implOptions{}, opts...) + + ctx = callbacks.EnsureRunInfo(ctx, r.GetType(), components.ComponentOfRetriever) + ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{ + Query: query, + TopK: *co.TopK, + Filter: io.FilterQuery, + ScoreThreshold: co.ScoreThreshold, + }) + defer func() { + if err != nil { + callbacks.OnError(ctx, err) + } + }() + + if io.FilterQuery != "" { + if err := validateFilterQuery(io.FilterQuery); err != nil { + return nil, err + } + } + + emb := co.Embedding + if emb == nil { + return nil, fmt.Errorf("[valkey retriever] embedding not provided") + } + + vectors, err := emb.EmbedStrings(r.makeEmbeddingCtx(ctx, emb), []string{query}) + if err != nil { + return nil, err + } + if len(vectors) != 1 { + return nil, fmt.Errorf("[valkey retriever] invalid return length of vector, got=%d, expected=1", len(vectors)) + } + + vecBytes := vector2Bytes(vectors[0]) + + var searchQuery string + params := []ftParam{ + {key: paramVector, value: string(vecBytes)}, + } + + if r.config.DistanceThreshold != nil { + params = append(params, ftParam{ + key: paramDistanceThreshold, + value: fmt.Sprintf("%f", *r.config.DistanceThreshold), + }) + baseQuery := fmt.Sprintf("@%s:[VECTOR_RANGE $%s $%s]", r.config.VectorField, paramDistanceThreshold, paramVector) + if io.FilterQuery != "" { + baseQuery = io.FilterQuery + " " + baseQuery + } + searchQuery = fmt.Sprintf("%s=>{$yield_distance_as: %s}", baseQuery, SortByDistanceAttributeName) + } else { + filter := "*" + if io.FilterQuery != "" { + filter = io.FilterQuery + } + searchQuery = fmt.Sprintf("(%s)=>[KNN %d @%s $%s AS %s]", + filter, *co.TopK, r.config.VectorField, paramVector, SortByDistanceAttributeName) + } + + // Build FT.SEARCH command args. + args := []string{"FT.SEARCH", *co.Index, searchQuery} + + // RETURN fields + args = append(args, "RETURN", strconv.Itoa(len(r.config.ReturnFields))) + args = append(args, r.config.ReturnFields...) + + // PARAMS + args = append(args, "PARAMS", strconv.Itoa(len(params)*2)) + for _, p := range params { + args = append(args, p.key, p.value) + } + + // LIMIT + args = append(args, "LIMIT", "0", strconv.Itoa(*co.TopK)) + + // DIALECT + args = append(args, "DIALECT", strconv.Itoa(r.config.Dialect)) + + raw, err := r.config.Client.CustomCommand(ctx, args) + if err != nil { + return nil, err + } + + searchDocs, err := parseFtSearchResponse(raw) + if err != nil { + return nil, err + } + + for _, sd := range searchDocs { + doc, convErr := r.config.DocumentConverter(ctx, sd) + if convErr != nil { + return nil, convErr + } + docs = append(docs, doc) + } + + callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: docs}) + return docs, nil +} + +func (r *Retriever) makeEmbeddingCtx(ctx context.Context, emb embedding.Embedder) context.Context { + runInfo := &callbacks.RunInfo{ + Component: components.ComponentOfEmbedding, + } + if embType, ok := components.GetType(emb); ok { + runInfo.Type = embType + } + runInfo.Name = runInfo.Type + string(runInfo.Component) + return callbacks.ReuseHandlers(ctx, runInfo) +} + +const typ = "Valkey" + +func (r *Retriever) GetType() string { + return typ +} + +func (r *Retriever) IsCallbacksEnabled() bool { + return true +} + +func defaultResultParser(returnFields []string) func(ctx context.Context, doc FtSearchDocument) (*schema.Document, error) { + return func(ctx context.Context, doc FtSearchDocument) (*schema.Document, error) { + resp := &schema.Document{ + ID: doc.Key, + Content: "", + MetaData: map[string]any{}, + } + for _, field := range returnFields { + val, found := doc.Fields[field] + if !found { + continue // skip missing fields gracefully + } + if field == defaultReturnFieldContent { + strVal, ok := val.(string) + if !ok { + return nil, fmt.Errorf("[defaultResultParser] field=%s: expected string, got %T", field, val) + } + resp.Content = strVal + } else if field == defaultReturnFieldVectorContent { + strVal, ok := val.(string) + if !ok { + return nil, fmt.Errorf("[defaultResultParser] field=%s: expected string, got %T", field, val) + } + resp.WithDenseVector(bytes2Vector([]byte(strVal))) + } else { + resp.MetaData[field] = val + } + } + return resp, nil + } +} + +type ftParam struct { + key string + value string +} + +// parseFtSearchResponse parses the raw CustomCommand response from FT.SEARCH. +// Expected format: []any{int64(totalCount), map[string]any{docKey: map[string]any{field: value}}} +func parseFtSearchResponse(raw any) ([]FtSearchDocument, error) { + arr, ok := raw.([]any) + if !ok || len(arr) < 2 { + return nil, nil + } + docMap, ok := arr[1].(map[string]any) + if !ok { + return nil, fmt.Errorf("[valkey retriever] unexpected FT.SEARCH response format") + } + docs := make([]FtSearchDocument, 0, len(docMap)) + for key, val := range docMap { + fields, ok := val.(map[string]any) + if !ok { + continue + } + docs = append(docs, FtSearchDocument{Key: key, Fields: fields}) + } + return docs, nil +} + +func vector2Bytes(vector []float64) []byte { + bytes := make([]byte, len(vector)*4) + for i, v := range vector { + binary.LittleEndian.PutUint32(bytes[i*4:], math.Float32bits(float32(v))) + } + return bytes +} + +func bytes2Vector(b []byte) []float64 { + n := len(b) / 4 + vector := make([]float64, n) + for i := 0; i < n; i++ { + bits := binary.LittleEndian.Uint32(b[i*4 : (i+1)*4]) + vector[i] = float64(math.Float32frombits(bits)) + } + return vector +} diff --git a/components/retriever/valkey/retriever_integration_test.go b/components/retriever/valkey/retriever_integration_test.go new file mode 100644 index 000000000..cf392f220 --- /dev/null +++ b/components/retriever/valkey/retriever_integration_test.go @@ -0,0 +1,267 @@ +//go:build integration + +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package valkey + +import ( + "context" + "net" + "os" + "strconv" + "testing" + "time" + + "github.com/cloudwego/eino/components/embedding" + glide "github.com/valkey-io/valkey-glide/go/v2" + "github.com/valkey-io/valkey-glide/go/v2/config" +) + +func getTestClient(t *testing.T) *glide.Client { + t.Helper() + addr := os.Getenv("VALKEY_ADDR") + if addr == "" { + addr = "localhost:6379" + } + host, portStr, err := net.SplitHostPort(addr) + if err != nil { + t.Fatalf("invalid VALKEY_ADDR %q: %v", addr, err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + t.Fatalf("invalid port in VALKEY_ADDR %q: %v", addr, err) + } + cfg := config.NewClientConfiguration(). + WithAddress(&config.NodeAddress{Host: host, Port: port}) + client, err := glide.NewClient(cfg) + if err != nil { + t.Fatalf("failed to create valkey client: %v", err) + } + return client +} + +// deterministicEmbedder returns vectors based on text for predictable similarity. +// Texts with similar characters produce similar vectors. +type deterministicEmbedder struct { + dim int +} + +func (e *deterministicEmbedder) EmbedStrings(_ context.Context, texts []string, _ ...embedding.Option) ([][]float64, error) { + vectors := make([][]float64, len(texts)) + for i, text := range texts { + vec := make([]float64, e.dim) + for j := 0; j < e.dim; j++ { + if j < len(text) { + vec[j] = float64(text[j]) / 255.0 + } + } + vectors[i] = vec + } + return vectors, nil +} + +func TestIntegration_Retriever_KNN(t *testing.T) { + ctx := context.Background() + client := getTestClient(t) + defer client.Close() + + const ( + prefix = "inttest:retriever:" + indexName = "inttest_retriever_idx" + dim = 8 + ) + emb := &deterministicEmbedder{dim: dim} + + // Cleanup + t.Cleanup(func() { + client.CustomCommand(ctx, []string{"FT.DROPINDEX", indexName}) + client.CustomCommand(ctx, []string{"DEL", prefix + "doc1", prefix + "doc2", prefix + "doc3"}) + }) + + // Drop index if it exists from a previous run + client.CustomCommand(ctx, []string{"FT.DROPINDEX", indexName}) + + // Create index + _, err := client.CustomCommand(ctx, []string{ + "FT.CREATE", indexName, + "ON", "HASH", + "PREFIX", "1", prefix, + "SCHEMA", + "content", "TEXT", + "vector_content", "VECTOR", "HNSW", "6", + "TYPE", "FLOAT32", + "DIM", "8", + "DISTANCE_METRIC", "L2", + }) + if err != nil { + t.Fatalf("FT.CREATE failed: %v", err) + } + + // Store documents using HSet directly (mimicking what the indexer does) + docs := []struct { + id string + content string + }{ + {"doc1", "Valkey is a high-performance key-value store"}, + {"doc2", "Vector search enables semantic similarity queries"}, + {"doc3", "The quick brown fox jumps over the lazy dog"}, + } + + for _, doc := range docs { + vecs, _ := emb.EmbedStrings(ctx, []string{doc.content}) + vecBytes := vector2Bytes(vecs[0]) + _, err := client.HSet(ctx, prefix+doc.id, map[string]string{ + "content": doc.content, + "vector_content": string(vecBytes), + }) + if err != nil { + t.Fatalf("HSet failed for %s: %v", doc.id, err) + } + } + + // Wait for index to be updated + time.Sleep(500 * time.Millisecond) + + // Create retriever + r, err := NewRetriever(ctx, &RetrieverConfig{ + Client: client, + Index: indexName, + TopK: 3, + Embedding: emb, + }) + if err != nil { + t.Fatalf("NewRetriever failed: %v", err) + } + + // Search - query similar to doc1 + results, err := r.Retrieve(ctx, "Valkey is a high-performance key-value store") + if err != nil { + t.Fatalf("Retrieve failed: %v", err) + } + + if len(results) == 0 { + t.Fatal("expected at least 1 result, got 0") + } + + t.Logf("Got %d results", len(results)) + for i, doc := range results { + t.Logf(" [%d] ID=%s Content=%q", i, doc.ID, doc.Content) + } + + // The first result should be doc1 since the query is identical + found := false + for _, doc := range results { + if doc.ID == prefix+"doc1" { + found = true + if doc.Content != "Valkey is a high-performance key-value store" { + t.Fatalf("unexpected content for doc1: %q", doc.Content) + } + break + } + } + if !found { + t.Fatal("expected doc1 in results (exact match query)") + } +} + +func TestIntegration_Retriever_VectorRange(t *testing.T) { + // VECTOR_RANGE is a Redis Stack feature not currently supported by Valkey Search. + // Valkey Search only supports KNN vector queries. + // This test is kept as a placeholder for when/if Valkey Search adds range support. + t.Skip("VECTOR_RANGE not supported by Valkey Search module") +} + +func TestIntegration_Retriever_WithFilter(t *testing.T) { + ctx := context.Background() + client := getTestClient(t) + defer client.Close() + + const ( + prefix = "inttest:filter:" + indexName = "inttest_filter_idx" + dim = 8 + ) + emb := &deterministicEmbedder{dim: dim} + + // Cleanup + t.Cleanup(func() { + client.CustomCommand(ctx, []string{"FT.DROPINDEX", indexName}) + client.CustomCommand(ctx, []string{"DEL", prefix + "doc1", prefix + "doc2"}) + }) + + client.CustomCommand(ctx, []string{"FT.DROPINDEX", indexName}) + + _, err := client.CustomCommand(ctx, []string{ + "FT.CREATE", indexName, + "ON", "HASH", + "PREFIX", "1", prefix, + "SCHEMA", + "content", "TEXT", + "category", "TAG", + "vector_content", "VECTOR", "HNSW", "6", + "TYPE", "FLOAT32", + "DIM", "8", + "DISTANCE_METRIC", "L2", + }) + if err != nil { + t.Fatalf("FT.CREATE failed: %v", err) + } + + // Store docs with different categories + vecs, _ := emb.EmbedStrings(ctx, []string{"hello world", "hello world"}) + + client.HSet(ctx, prefix+"doc1", map[string]string{ + "content": "hello world", + "category": "tech", + "vector_content": string(vector2Bytes(vecs[0])), + }) + client.HSet(ctx, prefix+"doc2", map[string]string{ + "content": "hello world", + "category": "science", + "vector_content": string(vector2Bytes(vecs[1])), + }) + + time.Sleep(500 * time.Millisecond) + + r, err := NewRetriever(ctx, &RetrieverConfig{ + Client: client, + Index: indexName, + TopK: 10, + Embedding: emb, + }) + if err != nil { + t.Fatalf("NewRetriever failed: %v", err) + } + + // Search with filter - only tech category + results, err := r.Retrieve(ctx, "hello world", WithFilterQuery("@category:{tech}")) + if err != nil { + t.Fatalf("Retrieve with filter failed: %v", err) + } + + t.Logf("Got %d filtered results", len(results)) + for i, doc := range results { + t.Logf(" [%d] ID=%s Content=%q", i, doc.ID, doc.Content) + } + + if len(results) != 1 { + t.Fatalf("expected 1 result with tech filter, got %d", len(results)) + } + if results[0].ID != prefix+"doc1" { + t.Fatalf("expected doc1 (tech), got %s", results[0].ID) + } +} diff --git a/components/retriever/valkey/retriever_test.go b/components/retriever/valkey/retriever_test.go new file mode 100644 index 000000000..a009fe59d --- /dev/null +++ b/components/retriever/valkey/retriever_test.go @@ -0,0 +1,399 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package valkey + +import ( + "context" + "fmt" + "testing" + + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/retriever" +) + +type mockSearchClient struct { + customCommandFn func(ctx context.Context, args []string) (any, error) +} + +func (m *mockSearchClient) CustomCommand(ctx context.Context, args []string) (any, error) { + if m.customCommandFn != nil { + return m.customCommandFn(ctx, args) + } + return nil, nil +} + +type mockEmbedding struct { + err error + cnt int + sizeForCall []int + dims int +} + +func (m *mockEmbedding) EmbedStrings(_ context.Context, _ []string, _ ...embedding.Option) ([][]float64, error) { + if m.err != nil { + return nil, m.err + } + if m.cnt >= len(m.sizeForCall) { + return nil, fmt.Errorf("unexpected call") + } + slice := make([]float64, m.dims) + for i := range slice { + slice[i] = 1.1 + } + r := make([][]float64, m.sizeForCall[m.cnt]) + m.cnt++ + for i := range r { + r[i] = slice + } + return r, nil +} + +func TestNewRetriever(t *testing.T) { + ctx := context.Background() + client := &mockSearchClient{} + + tests := []struct { + name string + config *RetrieverConfig + wantErr string + }{ + { + name: "embedding not provided", + config: &RetrieverConfig{Client: client, Index: "idx"}, + wantErr: "[NewRetriever] embedding not provided for valkey retriever", + }, + { + name: "index not provided", + config: &RetrieverConfig{Client: client, Embedding: &mockEmbedding{}}, + wantErr: "[NewRetriever] valkey index not provided", + }, + { + name: "client not provided", + config: &RetrieverConfig{Index: "idx", Embedding: &mockEmbedding{}}, + wantErr: "[NewRetriever] valkey client not provided", + }, + { + name: "success with defaults", + config: &RetrieverConfig{Client: client, Index: "idx", Embedding: &mockEmbedding{}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r, err := NewRetriever(ctx, tt.config) + if tt.wantErr != "" { + if err == nil || err.Error() != tt.wantErr { + t.Fatalf("expected error %q, got %v", tt.wantErr, err) + } + if r != nil { + t.Fatal("expected nil retriever on error") + } + } else { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if r == nil { + t.Fatal("expected non-nil retriever") + } + } + }) + } +} + +func TestNewRetriever_DoesNotMutateConfig(t *testing.T) { + ctx := context.Background() + cfg := &RetrieverConfig{ + Client: &mockSearchClient{}, + Index: "idx", + Embedding: &mockEmbedding{}, + } + _, err := NewRetriever(ctx, cfg) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.TopK != 0 { + t.Fatal("NewRetriever mutated caller's config") + } +} + +func TestRetriever_Retrieve_EmbeddingError(t *testing.T) { + ctx := context.Background() + client := &mockSearchClient{} + mockErr := fmt.Errorf("embedding failed") + + r := &Retriever{config: &RetrieverConfig{ + Client: client, + Index: "idx", + Embedding: &mockEmbedding{err: mockErr}, + TopK: 5, + VectorField: defaultReturnFieldVectorContent, + ReturnFields: []string{defaultReturnFieldContent, defaultReturnFieldVectorContent}, + DocumentConverter: defaultResultParser([]string{defaultReturnFieldContent, defaultReturnFieldVectorContent}), + Dialect: 2, + }} + + _, err := r.Retrieve(ctx, "test query") + if err == nil || err.Error() != "embedding failed" { + t.Fatalf("expected embedding error, got: %v", err) + } +} + +func TestRetriever_Retrieve_InvalidVectorLength(t *testing.T) { + ctx := context.Background() + client := &mockSearchClient{} + + r := &Retriever{config: &RetrieverConfig{ + Client: client, + Index: "idx", + Embedding: &mockEmbedding{sizeForCall: []int{2}, dims: 10}, + TopK: 5, + VectorField: defaultReturnFieldVectorContent, + ReturnFields: []string{defaultReturnFieldContent, defaultReturnFieldVectorContent}, + DocumentConverter: defaultResultParser([]string{defaultReturnFieldContent, defaultReturnFieldVectorContent}), + Dialect: 2, + }} + + _, err := r.Retrieve(ctx, "test query") + if err == nil { + t.Fatal("expected error for invalid vector length") + } + expected := "[valkey retriever] invalid return length of vector, got=2, expected=1" + if err.Error() != expected { + t.Fatalf("expected %q, got %q", expected, err.Error()) + } +} + +func TestRetriever_Retrieve_KNN(t *testing.T) { + ctx := context.Background() + expVec := make([]float64, 10) + for i := range expVec { + expVec[i] = 1.1 + } + + client := &mockSearchClient{ + customCommandFn: func(_ context.Context, args []string) (any, error) { + // Mock FT.SEARCH response: [totalCount, {key: map[field]value}] + return []any{ + int64(2), + map[string]any{ + "doc:1": map[string]any{"content": "hello", "vector_content": string(vector2Bytes(expVec))}, + "doc:2": map[string]any{"content": "world", "vector_content": string(vector2Bytes(expVec))}, + }, + }, nil + }, + } + + r, err := NewRetriever(ctx, &RetrieverConfig{ + Client: client, + Index: "test_index", + Embedding: &mockEmbedding{sizeForCall: []int{1}, dims: 10}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + docs, err := r.Retrieve(ctx, "test query") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(docs) != 2 { + t.Fatalf("expected 2 docs, got %d", len(docs)) + } + for _, doc := range docs { + if doc.Content != "hello" && doc.Content != "world" { + t.Fatalf("unexpected content: %s", doc.Content) + } + } +} + +func TestRetriever_Retrieve_VectorRange(t *testing.T) { + ctx := context.Background() + expVec := make([]float64, 10) + for i := range expVec { + expVec[i] = 1.1 + } + + client := &mockSearchClient{ + customCommandFn: func(_ context.Context, args []string) (any, error) { + return []any{ + int64(1), + map[string]any{ + "doc:1": map[string]any{"content": "hello", "vector_content": string(vector2Bytes(expVec))}, + }, + }, nil + }, + } + + dis := 10.0 + r, err := NewRetriever(ctx, &RetrieverConfig{ + Client: client, + Index: "test_index", + DistanceThreshold: &dis, + Embedding: &mockEmbedding{sizeForCall: []int{1}, dims: 10}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + docs, err := r.Retrieve(ctx, "test query") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(docs) != 1 { + t.Fatalf("expected 1 doc, got %d", len(docs)) + } + if docs[0].Content != "hello" { + t.Fatalf("unexpected content: %s", docs[0].Content) + } +} + +func TestVector2Bytes_Bytes2Vector_Roundtrip(t *testing.T) { + input := []float64{1.5, 2.5, 3.5, 0.0, -1.0} + b := vector2Bytes(input) + output := bytes2Vector(b) + if len(output) != len(input) { + t.Fatalf("length mismatch: got %d, want %d", len(output), len(input)) + } + for i := range input { + diff := input[i] - output[i] + if diff > 0.001 || diff < -0.001 { + t.Fatalf("value mismatch at %d: got %f, want %f", i, output[i], input[i]) + } + } +} + +func TestRetriever_GetType(t *testing.T) { + r := &Retriever{} + if r.GetType() != "Valkey" { + t.Fatalf("expected Valkey, got %s", r.GetType()) + } +} + +func TestRetriever_IsCallbacksEnabled(t *testing.T) { + r := &Retriever{} + if !r.IsCallbacksEnabled() { + t.Fatal("expected callbacks enabled") + } +} + +func TestDefaultResultParser(t *testing.T) { + parser := defaultResultParser([]string{defaultReturnFieldContent, defaultReturnFieldVectorContent}) + vec := []float64{1.0, 2.0} + + t.Run("success", func(t *testing.T) { + doc, err := parser(context.Background(), FtSearchDocument{ + Key: "id1", + Fields: map[string]any{ + defaultReturnFieldContent: "hello", + defaultReturnFieldVectorContent: string(vector2Bytes(vec)), + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if doc.ID != "id1" || doc.Content != "hello" { + t.Fatalf("unexpected doc: %+v", doc) + } + }) + + t.Run("missing field is skipped", func(t *testing.T) { + doc, err := parser(context.Background(), FtSearchDocument{ + Key: "id1", + Fields: map[string]any{defaultReturnFieldVectorContent: string(vector2Bytes(vec))}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if doc.Content != "" { + t.Fatalf("expected empty content for missing field, got %q", doc.Content) + } + }) + + t.Run("non-string content field returns error", func(t *testing.T) { + _, err := parser(context.Background(), FtSearchDocument{ + Key: "id1", + Fields: map[string]any{ + defaultReturnFieldContent: 123, + defaultReturnFieldVectorContent: string(vector2Bytes(vec)), + }, + }) + if err == nil { + t.Fatal("expected error for non-string content") + } + }) +} + +func TestRetriever_Retrieve_FilterQueryValidation(t *testing.T) { + ctx := context.Background() + + r := &Retriever{config: &RetrieverConfig{ + Client: &mockSearchClient{}, + Index: "idx", + Embedding: &mockEmbedding{sizeForCall: []int{1}, dims: 4}, + TopK: 5, + VectorField: defaultReturnFieldVectorContent, + ReturnFields: []string{defaultReturnFieldContent}, + DocumentConverter: defaultResultParser([]string{defaultReturnFieldContent}), + Dialect: 2, + }} + + tests := []struct { + name string + filter string + wantErr bool + }{ + {"valid tag filter", "@category:{electronics}", false}, + {"injection with =>", "@tag:{x}=>[KNN 100 @vec $v]", true}, + {"injection with KNN", "*)=>[KNN 9999 @vec $v AS d", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := r.Retrieve(ctx, "query", retriever.WrapImplSpecificOptFn(func(o *implOptions) { + o.FilterQuery = tt.filter + })) + if tt.wantErr && err == nil { + t.Fatal("expected error for disallowed filter") + } + if !tt.wantErr && err != nil && err.Error() == "[valkey retriever] filter contains disallowed syntax" { + t.Fatalf("unexpected validation error for valid filter: %v", err) + } + }) + } +} + +func TestValidateFilterQuery(t *testing.T) { + tests := []struct { + filter string + wantErr bool + }{ + {"@tag:{val}", false}, + {"@num:[0 100]", false}, + {"*", false}, + {"@x:{y}=>[KNN 10 @v $v]", true}, + {"something [KNN bypass", true}, + } + for _, tt := range tests { + err := validateFilterQuery(tt.filter) + if tt.wantErr && err == nil { + t.Errorf("expected error for %q", tt.filter) + } + if !tt.wantErr && err != nil { + t.Errorf("unexpected error for %q: %v", tt.filter, err) + } + } +}