add: retrieval timin

Author: iwilltry42

pkg/datastore/types/types.go

@@ -32,8 +32,13 @@ type Response struct {
 	ResultDocuments []vs.Document `json:"resultDocuments"`
 }
 
+type Stats struct {
+	RetrievalTimeSeconds float64 `json:"retrievalTimeSeconds,omitempty"`
+}
+
 type RetrievalResponse struct {
 	Query     string     `json:"originalQuery"`
 	Datasets  []string   `json:"queriedDatasets"`
 	Responses []Response `json:"subqueryResults"`
+	Stats     Stats      `json:"stats,omitempty"`
 }

pkg/flows/flows.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"log/slog"
 	"slices"
+	"time"
 
 	"github.com/acorn-io/z"
 	"github.com/google/uuid"
@@ -172,6 +173,8 @@ type RetrievalFlowOpts struct {
 }
 
 func (f *RetrievalFlow) Run(ctx context.Context, store store.Store, query string, datasetIDs []string, opts *RetrievalFlowOpts) (*dstypes.RetrievalResponse, error) {
+	retrievalFlowStartTime := time.Now()
+
 	if opts == nil {
 		opts = &RetrievalFlowOpts{}
 	}
@@ -213,5 +216,9 @@ func (f *RetrievalFlow) Run(ctx context.Context, store store.Store, query string
 	}
 	slog.Debug("Postprocessed RetrievalResponse", "num_responses", len(response.Responses), "original_query", query)
 
+	response.Stats = dstypes.Stats{
+		RetrievalTimeSeconds: time.Since(retrievalFlowStartTime).Seconds(),
+	}
+
 	return response, nil
 }

pkg/vectorstore/vectorstores.go

@@ -3,6 +3,7 @@ package vectorstore
 import (
 	"context"
 	"fmt"
+	"log/slog"
 	"strings"
 
 	etypes "github.com/gptscript-ai/knowledge/pkg/datastore/embeddings/types"
@@ -31,9 +32,9 @@ func New(ctx context.Context, dsn string, embeddingProvider etypes.EmbeddingMode
 		return nil, fmt.Errorf("failed to create embedding function: %w", err)
 	}
 
-	s := strings.Split(dsn, "://")
-	dialect := s[0]
-	dsn = s[1]
+	dialect := strings.Split(dsn, "://")[0]
+
+	slog.Debug("vectordb", "dialect", dialect, "dsn", dsn)
 
 	switch dialect {
 	case "chromem":