diff --git a/README.md b/README.md index af8fb2e..a962eb7 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ Because `chromem-go` is embeddable it enables you to add retrieval augmented gen It's *not* a library to connect to Chroma and also not a reimplementation of it in Go. It's a database on its own. -The focus is not scale (millions of documents) or number of features, but simplicity and performance for the most common use cases. On a mid-range 2020 Intel laptop CPU you can query 1,000 documents in 0.5 ms and 100,000 documents in 56 ms, both with just 44 memory allocations. See [Benchmarks](#benchmarks) for details. +The focus is not scale (millions of documents) or number of features, but simplicity and performance for the most common use cases. On a mid-range 2020 Intel laptop CPU you can query 1,000 documents in 0.5 ms and 100,000 documents in 50-60 ms, both with just 44 memory allocations. See [Benchmarks](#benchmarks) for details. > ⚠️ The project is in beta, under heavy construction, and may introduce breaking changes in releases before `v1.0.0`. All changes are documented in the [`CHANGELOG`](./CHANGELOG.md). @@ -183,13 +183,13 @@ For full, working examples, using the vector database for retrieval augmented ge ## Benchmarks -Benchmarked on 2024-03-16 with: +Benchmarked on 2024-03-17 with: - Computer: Framework Laptop 13 (first generation, 2021) - CPU: 11th Gen Intel Core i5-1135G7 (2020) - Memory: 32 GB - OS: Fedora Linux 39 -- Kernel: 6.7 + - Kernel: 6.7 ```console $ go test -benchmem -run=^$ -bench . @@ -197,18 +197,18 @@ goos: linux goarch: amd64 pkg: github.com/philippgille/chromem-go cpu: 11th Gen Intel(R) Core(TM) i5-1135G7 @ 2.40GHz -BenchmarkCollection_Query_NoContent_100-8 10000 110126 ns/op 6492 B/op 44 allocs/op -BenchmarkCollection_Query_NoContent_1000-8 2020 537416 ns/op 35669 B/op 44 allocs/op -BenchmarkCollection_Query_NoContent_5000-8 351 4264192 ns/op 166728 B/op 44 allocs/op -BenchmarkCollection_Query_NoContent_25000-8 75 16411744 ns/op 813928 B/op 44 allocs/op -BenchmarkCollection_Query_NoContent_100000-8 18 64670962 ns/op 3205962 B/op 44 allocs/op -BenchmarkCollection_Query_100-8 10923 109936 ns/op 6480 B/op 44 allocs/op -BenchmarkCollection_Query_1000-8 2184 562778 ns/op 35667 B/op 44 allocs/op -BenchmarkCollection_Query_5000-8 400 2986732 ns/op 166750 B/op 44 allocs/op -BenchmarkCollection_Query_25000-8 88 15433911 ns/op 813896 B/op 44 allocs/op -BenchmarkCollection_Query_100000-8 19 63696478 ns/op 3205982 B/op 44 allocs/op +BenchmarkCollection_Query_NoContent_100-8 10000 109957 ns/op 6487 B/op 44 allocs/op +BenchmarkCollection_Query_NoContent_1000-8 2043 541337 ns/op 35667 B/op 44 allocs/op +BenchmarkCollection_Query_NoContent_5000-8 361 4230542 ns/op 166729 B/op 44 allocs/op +BenchmarkCollection_Query_NoContent_25000-8 79 16343977 ns/op 813902 B/op 44 allocs/op +BenchmarkCollection_Query_NoContent_100000-8 18 63417540 ns/op 3205961 B/op 44 allocs/op +BenchmarkCollection_Query_100-8 10861 110167 ns/op 6480 B/op 44 allocs/op +BenchmarkCollection_Query_1000-8 2146 540489 ns/op 35669 B/op 44 allocs/op +BenchmarkCollection_Query_5000-8 442 4913970 ns/op 166748 B/op 44 allocs/op +BenchmarkCollection_Query_25000-8 88 15213089 ns/op 813909 B/op 44 allocs/op +BenchmarkCollection_Query_100000-8 19 55963675 ns/op 3206027 B/op 44 allocs/op PASS -ok github.com/philippgille/chromem-go 31.373s +ok github.com/philippgille/chromem-go 27.887s ``` ## Motivation diff --git a/collection.go b/collection.go index 92454ff..fda0fe5 100644 --- a/collection.go +++ b/collection.go @@ -266,7 +266,8 @@ type Result struct { // Performs an exhaustive nearest neighbor search on the collection. // -// - queryText: The text to search for. +// - queryText: The text to search for. Its embedding will be created using the +// collection's embedding function. // - nResults: The number of results to return. Must be > 0. // - where: Conditional filtering on metadata. Optional. // - whereDocument: Conditional filtering on documents. Optional. @@ -274,6 +275,26 @@ func (c *Collection) Query(ctx context.Context, queryText string, nResults int, if queryText == "" { return nil, errors.New("queryText is empty") } + + queryVectors, err := c.embed(ctx, queryText) + if err != nil { + return nil, fmt.Errorf("couldn't create embedding of query: %w", err) + } + + return c.QueryEmbedding(ctx, queryVectors, nResults, where, whereDocument) +} + +// Performs an exhaustive nearest neighbor search on the collection. +// +// - queryEmbedding: The embedding of the query to search for. It must be created +// with the same embedding model as the document embeddings in the collection. +// - nResults: The number of results to return. Must be > 0. +// - where: Conditional filtering on metadata. Optional. +// - whereDocument: Conditional filtering on documents. Optional. +func (c *Collection) QueryEmbedding(ctx context.Context, queryEmbedding []float32, nResults int, where, whereDocument map[string]string) ([]Result, error) { + if len(queryEmbedding) == 0 { + return nil, errors.New("queryEmbedding is empty") + } if nResults <= 0 { return nil, errors.New("nResults must be > 0") } @@ -302,13 +323,8 @@ func (c *Collection) Query(ctx context.Context, queryText string, nResults int, return nil, nil } - queryVectors, err := c.embed(ctx, queryText) - if err != nil { - return nil, fmt.Errorf("couldn't create embedding of query: %w", err) - } - // For the remaining documents, calculate cosine similarity. - docSim, err := calcDocSimilarity(ctx, queryVectors, filteredDocs) + docSim, err := calcDocSimilarity(ctx, queryEmbedding, filteredDocs) if err != nil { return nil, fmt.Errorf("couldn't calculate cosine similarity: %w", err) } diff --git a/collection_test.go b/collection_test.go index 89ffd74..807b4f6 100644 --- a/collection_test.go +++ b/collection_test.go @@ -477,16 +477,13 @@ func benchmarkCollection_Query(b *testing.B, n int, withContent bool) { } // Most embeddings are normalized, so we normalize this one too qv = normalizeVector(qv) - embeddingFunc := func(_ context.Context, text string) ([]float32, error) { - if text != "foo" { - return nil, errors.New("embedding func not expected to be called") - } - return qv, nil - } // Create collection db := NewDB() name := "test" + embeddingFunc := func(_ context.Context, text string) ([]float32, error) { + return nil, errors.New("embedding func not expected to be called") + } c, err := db.CreateCollection(name, nil, embeddingFunc) if err != nil { b.Fatal("expected no error, got", err) @@ -524,7 +521,7 @@ func benchmarkCollection_Query(b *testing.B, n int, withContent bool) { // Query var res []Result for i := 0; i < b.N; i++ { - res, err = c.Query(ctx, "foo", 10, nil, nil) + res, err = c.QueryEmbedding(ctx, qv, 10, nil, nil) } if err != nil { b.Fatal("expected nil, got", err)