diff --git a/README.md b/README.md index 50533b6..d13a4a0 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,7 @@ New to Qdrant? Here's how to get started: | --- | --- | --- | --- | | ✅ | **Chonkie** | Advanced text chunking with Qdrant handshake | [Tutorial](chonkie/Chonkie_Qdrant_Handshake.ipynb) | | 🚧 | **FastEmbed** | Fast, lightweight embedding library | Coming Soon | +| ✅ | **bm25-go-hybrid** | Dense + sparse (BM25) hybrid search in Go, using a FastEmbed-compatible sparse encoder | [Demo](bm25-go-hybrid/) | | 🚧 | **LangChain** | Python framework for LLM applications | Coming Soon | | 🚧 | **LlamaIndex** | Data framework for LLM applications | Coming Soon | diff --git a/bm25-go-hybrid/README.md b/bm25-go-hybrid/README.md new file mode 100644 index 0000000..f7d35cf --- /dev/null +++ b/bm25-go-hybrid/README.md @@ -0,0 +1,48 @@ +# BM25 hybrid search in Go + +A runnable demo of **dense + sparse hybrid search** against Qdrant from Go. + +- **Sparse (BM25)** vectors come from + [`github.com/harsh04/bm25`](https://github.com/harsh04/bm25), a byte-for-byte Go + port of FastEmbed's `Qdrant/bm25` encoder — the same sparse vectors you'd get + from Python FastEmbed, produced client-side in Go. It emits term frequencies + only; Qdrant applies IDF server-side via the sparse vector's modifier. +- **Dense** vectors come from Google's Gemini embedding API (`gemini-embedding-001`). +- The two arms are fused with **Reciprocal Rank Fusion (RRF)** by Qdrant. + +## Run + +```sh +export GEMINI_API_KEY=... # https://aistudio.google.com/apikey +docker run -p 6333:6333 -p 6334:6334 qdrant/qdrant # any Qdrant works +go run . +``` + +## Expected output + +The first query overlaps the docs word-for-word, so BM25 alone would find it. The +other two share **no words** with the doc they should match — the dense arm is +what surfaces them. That's the whole point of hybrid: + +``` +query: "when can I check in" + 1.0000 Check-in is at 3 PM and check-out is at 11 AM. + 0.3333 Breakfast is served daily from 7 to 10. + 0.2500 The swimming pool is open from 6 AM to 9 PM. + +query: "what time should I arrive" + 0.5000 Check-in is at 3 PM and check-out is at 11 AM. + 0.3333 The swimming pool is open from 6 AM to 9 PM. + 0.2500 Breakfast is served daily from 7 to 10. + +query: "somewhere to eat early" + 0.5000 Breakfast is served daily from 7 to 10. + 0.3333 Check-in is at 3 PM and check-out is at 11 AM. + 0.2500 The swimming pool is open from 6 AM to 9 PM. +``` + +(Scores are RRF ranks, so they're stable; the top hit for each query is the +semantically correct doc.) + +> Prefer less wiring? The [`qhybrid`](https://pkg.go.dev/github.com/harsh04/bm25/qhybrid) +> module wraps collection setup, upsert, and this RRF query behind a small API. diff --git a/bm25-go-hybrid/go.mod b/bm25-go-hybrid/go.mod new file mode 100644 index 0000000..847c725 --- /dev/null +++ b/bm25-go-hybrid/go.mod @@ -0,0 +1,19 @@ +module bm25-go-hybrid-demo + +go 1.25.0 + +require ( + github.com/harsh04/bm25 v1.3.2 + github.com/qdrant/go-client v1.18.2 +) + +require ( + github.com/blevesearch/snowballstem v0.9.0 // indirect + github.com/twmb/murmur3 v1.1.8 // indirect + golang.org/x/net v0.53.0 // indirect + golang.org/x/sys v0.43.0 // indirect + golang.org/x/text v0.36.0 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260427160629-7cedc36a6bc4 // indirect + google.golang.org/grpc v1.80.0 // indirect + google.golang.org/protobuf v1.36.11 // indirect +) diff --git a/bm25-go-hybrid/go.sum b/bm25-go-hybrid/go.sum new file mode 100644 index 0000000..01c5b08 --- /dev/null +++ b/bm25-go-hybrid/go.sum @@ -0,0 +1,20 @@ +github.com/blevesearch/snowballstem v0.9.0 h1:lMQ189YspGP6sXvZQ4WZ+MLawfV8wOmPoD/iWeNXm8s= +github.com/blevesearch/snowballstem v0.9.0/go.mod h1:PivSj3JMc8WuaFkTSRDW2SlrulNWPl4ABg1tC/hlgLs= +github.com/harsh04/bm25 v1.3.2 h1:0XMb+cYGq1FhULWWDIkcT1ntnpgHGWHU8YHeJPuwruo= +github.com/harsh04/bm25 v1.3.2/go.mod h1:V5VS1CDdUWM8Y9sjxPS14t341RJLGlcNmJYF1xZF/Vo= +github.com/qdrant/go-client v1.18.2 h1:7ViiXB/fB4vfzdUtEYZ7g2vSG+yMU6EMu8CaNWF1g1c= +github.com/qdrant/go-client v1.18.2/go.mod h1:Xkfp+r89uNOgSbvilVAhCZ3wKI4G+hB/r9Zr2m4zifI= +github.com/twmb/murmur3 v1.1.8 h1:8Yt9taO/WN3l08xErzjeschgZU2QSrwm1kclYq+0aRg= +github.com/twmb/murmur3 v1.1.8/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ= +golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA= +golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs= +golang.org/x/sys v0.43.0 h1:Rlag2XtaFTxp19wS8MXlJwTvoh8ArU6ezoyFsMyCTNI= +golang.org/x/sys v0.43.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.36.0 h1:JfKh3XmcRPqZPKevfXVpI1wXPTqbkE5f7JA92a55Yxg= +golang.org/x/text v0.36.0/go.mod h1:NIdBknypM8iqVmPiuco0Dh6P5Jcdk8lJL0CUebqK164= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260427160629-7cedc36a6bc4 h1:tEkOQcXgF6dH1G+MVKZrfpYvozGrzb91k6ha7jireSM= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260427160629-7cedc36a6bc4/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.80.0 h1:Xr6m2WmWZLETvUNvIUmeD5OAagMw3FiKmMlTdViWsHM= +google.golang.org/grpc v1.80.0/go.mod h1:ho/dLnxwi3EDJA4Zghp7k2Ec1+c2jqup0bFkw07bwF4= +google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= +google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= diff --git a/bm25-go-hybrid/main.go b/bm25-go-hybrid/main.go new file mode 100644 index 0000000..9640b48 --- /dev/null +++ b/bm25-go-hybrid/main.go @@ -0,0 +1,171 @@ +// Demo: dense + sparse (BM25) hybrid search against Qdrant, in Go. +// +// Sparse vectors come from github.com/harsh04/bm25 (FastEmbed Qdrant/bm25 +// parity). Dense vectors come from Google's Gemini embedding API. The two arms +// are fused with Reciprocal Rank Fusion (RRF) server-side by Qdrant. +// +// Run: +// +// export GEMINI_API_KEY=... # https://aistudio.google.com/apikey +// docker run -p 6333:6333 -p 6334:6334 qdrant/qdrant +// go run . +package main + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "math" + "net/http" + "os" + + "github.com/harsh04/bm25" + "github.com/qdrant/go-client/qdrant" +) + +const ( + collection = "bm25_go_hybrid_demo" + embedModel = "gemini-embedding-001" + denseDim = 1536 // Gemini embedding output dimensionality +) + +func ptr[T any](v T) *T { return &v } + +// denseEmbed returns a real semantic embedding from the Gemini API. Truncated +// (sub-3072) outputs are L2-normalized, as Google recommends, so cosine behaves. +func denseEmbed(ctx context.Context, text string) ([]float32, error) { + key := os.Getenv("GEMINI_API_KEY") + if key == "" { + return nil, fmt.Errorf("set GEMINI_API_KEY (https://aistudio.google.com/apikey)") + } + reqBody, _ := json.Marshal(map[string]any{ + "content": map[string]any{"parts": []map[string]string{{"text": text}}}, + "outputDimensionality": denseDim, + }) + url := "https://generativelanguage.googleapis.com/v1beta/models/" + embedModel + ":embedContent" + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(reqBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-goog-api-key", key) // header, not URL, so it can't leak into errors + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("gemini %d: %s", resp.StatusCode, b) + } + var out struct { + Embedding struct { + Values []float32 `json:"values"` + } `json:"embedding"` + } + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + return nil, err + } + return l2norm(out.Embedding.Values), nil +} + +func l2norm(v []float32) []float32 { + var sum float64 + for _, x := range v { + sum += float64(x) * float64(x) + } + if sum == 0 { + return v + } + n := float32(math.Sqrt(sum)) + for i := range v { + v[i] /= n + } + return v +} + +func main() { + ctx := context.Background() + client, err := qdrant.NewClient(&qdrant.Config{Host: "127.0.0.1", Port: 6334}) + if err != nil { + log.Fatal(err) + } + enc := bm25.New() + + // Collection: a dense vector + a sparse slot with the IDF modifier. + _ = client.DeleteCollection(ctx, collection) + if err := client.CreateCollection(ctx, &qdrant.CreateCollection{ + CollectionName: collection, + VectorsConfig: qdrant.NewVectorsConfigMap(map[string]*qdrant.VectorParams{ + "dense": {Size: denseDim, Distance: qdrant.Distance_Cosine}, + }), + SparseVectorsConfig: qdrant.NewSparseVectorsConfig(map[string]*qdrant.SparseVectorParams{ + "sparse": {Modifier: qdrant.Modifier_Idf.Enum()}, + }), + }); err != nil { + log.Fatal(err) + } + + docs := []string{ + "Check-in is at 3 PM and check-out is at 11 AM.", + "Free WiFi is available throughout the hotel.", + "The swimming pool is open from 6 AM to 9 PM.", + "Pets are welcome with a cleaning fee.", + "Breakfast is served daily from 7 to 10.", + "The airport shuttle runs every 30 minutes.", + } + var points []*qdrant.PointStruct + for i, d := range docs { + dense, err := denseEmbed(ctx, d) + if err != nil { + log.Fatal(err) + } + idx, val := enc.Encode(d) // sparse BM25 vector + points = append(points, &qdrant.PointStruct{ + Id: qdrant.NewIDNum(uint64(i)), + Vectors: qdrant.NewVectorsMap(map[string]*qdrant.Vector{ + "dense": qdrant.NewVector(dense...), + "sparse": qdrant.NewVectorSparse(idx, val), + }), + Payload: qdrant.NewValueMap(map[string]any{"text": d}), + }) + } + if _, err := client.Upsert(ctx, &qdrant.UpsertPoints{ + CollectionName: collection, Points: points, Wait: ptr(true), + }); err != nil { + log.Fatal(err) + } + + // The first query shares words with a doc, so BM25 alone finds it. The other + // two use different words from the docs they should match, so the dense arm + // is what surfaces them — that's the point of hybrid. + queries := []string{ + "when can I check in", // lexical overlap -> sparse (BM25) finds it + "what time should I arrive", // same meaning, no shared words -> dense arm + "somewhere to eat early", // -> breakfast, purely semantic -> dense arm + } + for _, query := range queries { + dense, err := denseEmbed(ctx, query) + if err != nil { + log.Fatal(err) + } + sIdx, sVal := enc.Encode(query) + hits, err := client.Query(ctx, &qdrant.QueryPoints{ + CollectionName: collection, + Prefetch: []*qdrant.PrefetchQuery{ + {Query: qdrant.NewQueryDense(dense), Using: ptr("dense"), Limit: ptr(uint64(10))}, + {Query: qdrant.NewQuerySparse(sIdx, sVal), Using: ptr("sparse"), Limit: ptr(uint64(10))}, + }, + Query: qdrant.NewQueryFusion(qdrant.Fusion_RRF), // fuse dense + sparse + Limit: ptr(uint64(3)), + WithPayload: qdrant.NewWithPayload(true), + }) + if err != nil { + log.Fatal(err) + } + fmt.Printf("\nquery: %q\n", query) + for _, h := range hits { + fmt.Printf(" %.4f %s\n", h.Score, h.Payload["text"].GetStringValue()) + } + } +}