Merge pull request #13 from jhrozek/options

Add options to QueryRelevantDocuments, add a new WithScoreThreshold option

Author: lukehinds

examples/qdrant/main.go

@@ -72,7 +72,9 @@ func main() {
 	}
 
 	// Query the most relevant documents based on a given embedding
-	retrievedDocs, err := vectorDB.QueryRelevantDocuments(ctx, queryEmbedding, 5, collection_name)
+	retrievedDocs, err := vectorDB.QueryRelevantDocuments(
+		ctx, queryEmbedding, collection_name,
+		db.WithLimit(5), db.WithScoreThreshold(0.7))
 	if err != nil {
 		log.Fatalf("Failed to query documents: %v", err)
 	}

pkg/db/qdrant.go

@@ -88,25 +88,48 @@ func (qv *QdrantVector) SaveEmbeddings(ctx context.Context, docID string, embedd
 	return nil
 }
 
+// QueryOpt represents an option for a query. This is the type that should
+// be returned from query options functions.
+type QueryOpt func(*qdrant.QueryPoints)
+
+// WithLimit sets the limit of the number of documents to return in a query.
+func WithLimit(limit uint64) QueryOpt {
+	return func(q *qdrant.QueryPoints) {
+		q.Limit = &limit
+	}
+}
+
+// WithScoreThreshold sets the score threshold for a query. The higher the threshold, the more relevant the results.
+func WithScoreThreshold(threshold float32) QueryOpt {
+	return func(q *qdrant.QueryPoints) {
+		q.ScoreThreshold = &threshold
+	}
+}
+
 // QueryRelevantDocuments retrieves the most relevant documents based on a given embedding.
 //
 // Parameters:
 //   - ctx: The context for the query.
 //   - embedding: The query embedding.
 //   - limit: The number of documents to return.
+//   - collection: The collection name to query.
 //
 // Returns:
 //   - A slice of QDrantDocument structs containing the most relevant documents.
 //   - An error if the query fails.
-func (qv *QdrantVector) QueryRelevantDocuments(ctx context.Context, embedding []float32, limit int, colllection string) ([]Document, error) {
-	limitUint := uint64(limit) // Convert limit to uint64
+func (qv *QdrantVector) QueryRelevantDocuments(
+	ctx context.Context, embedding []float32, collection string, queryOpts ...QueryOpt,
+) ([]Document, error) {
 	query := &qdrant.QueryPoints{
-		CollectionName: colllection, // Replace with actual collection name
+		CollectionName: collection, // Replace with actual collection name
 		Query:          qdrant.NewQuery(embedding...),
-		Limit:          &limitUint,
 		WithPayload:    qdrant.NewWithPayloadInclude("content"),
 	}
 
+	for _, opt := range queryOpts {
+		opt(query)
+	}
+
 	response, err := qv.client.Query(ctx, query)
 	if err != nil {
 		return nil, fmt.Errorf("failed to search points: %w", err)

pkg/db/qdrant_test.go

@@ -84,15 +84,19 @@ func (t *testQdrantVector) SaveEmbeddings(ctx context.Context, docID string, emb
 	return err
 }
 
-func (t *testQdrantVector) QueryRelevantDocuments(ctx context.Context, embedding []float32, limit int, collection string) ([]Document, error) {
-	limitUint := uint64(limit)
+func (t *testQdrantVector) QueryRelevantDocuments(
+	ctx context.Context, embedding []float32, collection string, queryOpts ...QueryOpt,
+) ([]Document, error) {
 	query := &qdrant.QueryPoints{
 		CollectionName: collection,
 		Query:          qdrant.NewQuery(embedding...),
-		Limit:          &limitUint,
 		WithPayload:    qdrant.NewWithPayloadInclude("content"),
 	}
 
+	for _, opt := range queryOpts {
+		opt(query)
+	}
+
 	response, err := t.mockClient.Query(ctx, query)
 	if err != nil {
 		return nil, err
@@ -186,7 +190,8 @@ func TestQueryRelevantDocuments(t *testing.T) {
 	})).Return(mockResponse, nil)
 
 	// Test the QueryRelevantDocuments function
-	docs, err := qv.QueryRelevantDocuments(ctx, embedding, limit, collection)
+	docs, err := qv.QueryRelevantDocuments(ctx, embedding, collection,
+		WithLimit(5), WithScoreThreshold(0.7))
 	assert.NoError(t, err)
 	assert.Len(t, docs, 1)
 	assert.Equal(t, "test content", docs[0].Metadata["content"])