add: get and list dataset cmds

pkg/cmd/get_dataset.go

@@ -0,0 +1,44 @@
+package cmd
+
+import (
+	"encoding/json"
+	"fmt"
+	"github.com/spf13/cobra"
+)
+
+type ClientGetDataset struct {
+	Client
+}
+
+func (s *ClientGetDataset) Customize(cmd *cobra.Command) {
+	cmd.Use = "get-dataset <dataset-id>"
+	cmd.Short = "Get a dataset"
+	cmd.Args = cobra.ExactArgs(1)
+}
+
+func (s *ClientGetDataset) Run(cmd *cobra.Command, args []string) error {
+	c, err := s.getClient()
+	if err != nil {
+		return err
+	}
+
+	datasetID := args[0]
+
+	ds, err := c.GetDataset(cmd.Context(), datasetID)
+	if err != nil {
+		return fmt.Errorf("failed to get dataset: %w", err)
+	}
+
+	if ds == nil {
+		fmt.Println("dataset not found")
+		return fmt.Errorf("dataset not found")
+	}
+
+	jsonOutput, err := json.Marshal(ds)
+	if err != nil {
+		return fmt.Errorf("failed to marshal dataset: %w", err)
+	}
+
+	fmt.Println(string(jsonOutput))
+	return nil
+}

pkg/cmd/list_datasets.go

@@ -0,0 +1,42 @@
+package cmd
+
+import (
+	"encoding/json"
+	"fmt"
+	"github.com/spf13/cobra"
+)
+
+type ClientListDatasets struct {
+	Client
+}
+
+func (s *ClientListDatasets) Customize(cmd *cobra.Command) {
+	cmd.Use = "list-datasets"
+	cmd.Short = "List existing datasets"
+	cmd.Args = cobra.NoArgs
+}
+
+func (s *ClientListDatasets) Run(cmd *cobra.Command, args []string) error {
+	c, err := s.getClient()
+	if err != nil {
+		return err
+	}
+
+	ds, err := c.ListDatasets(cmd.Context())
+	if err != nil {
+		return fmt.Errorf("failed to list datasets: %w", err)
+	}
+
+	if len(ds) == 0 {
+		fmt.Println("no datasets found")
+		return nil
+	}
+
+	jsonOutput, err := json.Marshal(ds)
+	if err != nil {
+		return fmt.Errorf("failed to marshal datasets: %w", err)
+	}
+
+	fmt.Println(string(jsonOutput))
+	return nil
+}

pkg/cmd/main.go

@@ -6,7 +6,16 @@ import (
 )
 
 func New() *cobra.Command {
-	return cmd.Command(&Knowledge{}, new(Server), new(ClientCreateDataset), new(ClientIngest), new(ClientDeleteDataset), new(ClientRetrieve))
+	return cmd.Command(
+		&Knowledge{},
+		new(Server),
+		new(ClientCreateDataset),
+		new(ClientGetDataset),
+		new(ClientListDatasets),
+		new(ClientIngest),
+		new(ClientDeleteDataset),
+		new(ClientRetrieve),
+	)
 }
 
 type Knowledge struct{}

pkg/datastore/dataset.go

@@ -3,6 +3,7 @@ package datastore
 import (
 	"context"
 	"errors"
+	"fmt"
 	"github.com/acorn-io/z"
 	"github.com/gptscript-ai/knowledge/pkg/index"
 	"github.com/gptscript-ai/knowledge/pkg/types"
@@ -50,13 +51,13 @@ func (s *Datastore) DeleteDataset(ctx context.Context, datasetID string) error {
 
 func (s *Datastore) GetDataset(ctx context.Context, datasetID string) (*index.Dataset, error) {
 	// Get dataset with files and associated documents preloaded
-	var dataset *index.Dataset
+	dataset := &index.Dataset{}
 	tx := s.Index.WithContext(ctx).Preload("Files.Documents").First(dataset, "id = ?", datasetID)
 	if tx.Error != nil {
 		if errors.Is(tx.Error, gorm.ErrRecordNotFound) {
 			return nil, nil
 		}
-		return nil, tx.Error
+		return nil, fmt.Errorf("failed to get dataset %q from DB: %w", datasetID, tx.Error)
 	}
 
 	return dataset, nil

pkg/datastore/document.go

@@ -7,9 +7,6 @@ import (
 	"log/slog"
 )
 
-// ErrDBDocumentNotFound is returned when a document is not found in the database.
-var ErrDBDocumentNotFound = fmt.Errorf("document not found in database")
-
 func (s *Datastore) DeleteDocument(ctx context.Context, documentID, datasetID string) error {
 	// Find in Database
 	var document index.Document

pkg/datastore/errors.go

@@ -0,0 +1,8 @@
+package datastore
+
+import (
+	"errors"
+)
+
+// ErrDBDocumentNotFound is returned when a document is not found in the database.
+var ErrDBDocumentNotFound = errors.New("document not found in database")