Jobs Indexing Operations (#1754)

* vendor changes

* added changes for indexing operations

* generated mocks

* reversed modules.txt change

* fixed integration test

* added other change

Author: DO-rrao

commands/displayers/genai.go

@@ -201,6 +201,165 @@ func (v *KnowledgeBaseDataSource) KV() []map[string]any {
 	return out
 }
 
+// IndexingJobDataSource displayer
+type IndexingJobDataSource struct {
+	IndexingJobDataSources do.IndexingJobDataSources
+}
+
+var _ Displayable = &IndexingJobDataSource{}
+
+func (v *IndexingJobDataSource) JSON(out io.Writer) error {
+	return writeJSON(v.IndexingJobDataSources, out)
+}
+
+func (v *IndexingJobDataSource) ColMap() map[string]string {
+	return map[string]string{
+		"DataSourceUuid":    "Data Source UUID",
+		"Status":            "Status",
+		"StartedAt":         "Started At",
+		"CompletedAt":       "Completed At",
+		"IndexedItemCount":  "Indexed Items",
+		"FailedItemCount":   "Failed Items",
+		"SkippedItemCount":  "Skipped Items",
+		"RemovedItemCount":  "Removed Items",
+		"IndexedFileCount":  "Indexed Files",
+		"TotalFileCount":    "Total Files",
+		"TotalBytes":        "Total Bytes",
+		"TotalBytesIndexed": "Total Bytes Indexed",
+		"ErrorMsg":          "Error Message",
+		"ErrorDetails":      "Error Details",
+	}
+}
+
+func (v *IndexingJobDataSource) Cols() []string {
+	return []string{
+		"DataSourceUuid",
+		"Status",
+		"StartedAt",
+		"CompletedAt",
+		"IndexedItemCount",
+		"FailedItemCount",
+		"SkippedItemCount",
+		"IndexedFileCount",
+		"TotalFileCount",
+	}
+}
+
+func (v *IndexingJobDataSource) KV() []map[string]any {
+	if v == nil || v.IndexingJobDataSources == nil {
+		return []map[string]any{}
+	}
+	out := make([]map[string]any, 0, len(v.IndexingJobDataSources))
+	for _, ds := range v.IndexingJobDataSources {
+		startedAt := ""
+		if ds.StartedAt != nil {
+			startedAt = ds.StartedAt.String()
+		}
+		completedAt := ""
+		if ds.CompletedAt != nil {
+			completedAt = ds.CompletedAt.String()
+		}
+
+		out = append(out, map[string]any{
+			"DataSourceUuid":    ds.DataSourceUuid,
+			"Status":            ds.Status,
+			"StartedAt":         startedAt,
+			"CompletedAt":       completedAt,
+			"IndexedItemCount":  ds.IndexedItemCount,
+			"FailedItemCount":   ds.FailedItemCount,
+			"SkippedItemCount":  ds.SkippedItemCount,
+			"RemovedItemCount":  ds.RemovedItemCount,
+			"IndexedFileCount":  ds.IndexedFileCount,
+			"TotalFileCount":    ds.TotalFileCount,
+			"TotalBytes":        ds.TotalBytes,
+			"TotalBytesIndexed": ds.TotalBytesIndexed,
+			"ErrorMsg":          ds.ErrorMsg,
+			"ErrorDetails":      ds.ErrorDetails,
+		})
+	}
+	return out
+}
+
+// IndexingJob displayer
+type IndexingJob struct {
+	IndexingJobs do.IndexingJobs
+}
+
+var _ Displayable = &IndexingJob{}
+
+func (v *IndexingJob) JSON(out io.Writer) error {
+	return writeJSON(v.IndexingJobs, out)
+}
+
+func (v *IndexingJob) ColMap() map[string]string {
+	return map[string]string{
+		"CompletedDatasources": "Completed Datasources",
+		"CreatedAt":            "Created At",
+		"DataSourceUuids":      "Data Source UUIDs",
+		"FinishedAt":           "Finished At",
+		"KnowledgeBaseUuid":    "Knowledge Base UUID",
+		"Phase":                "Phase",
+		"StartedAt":            "Started At",
+		"Status":               "Status",
+		"Tokens":               "Tokens",
+		"TotalDatasources":     "Total Datasources",
+		"TotalItemsFailed":     "Total Items Failed",
+		"TotalItemsIndexed":    "Total Items Indexed",
+		"TotalItemsSkipped":    "Total Items Skipped",
+		"UpdatedAt":            "Updated At",
+		"UUID":                 "UUID",
+	}
+}
+
+func (v *IndexingJob) Cols() []string {
+	return []string{
+		"UUID",
+		"KnowledgeBaseUuid",
+		"Phase",
+		"Status",
+		"CompletedDatasources",
+		"TotalDatasources",
+		"Tokens",
+		"TotalItemsIndexed",
+		"TotalItemsFailed",
+		"TotalItemsSkipped",
+		"CreatedAt",
+		"StartedAt",
+		"FinishedAt",
+		"UpdatedAt",
+	}
+}
+
+func (v *IndexingJob) KV() []map[string]any {
+	if v == nil || v.IndexingJobs == nil {
+		return []map[string]any{}
+	}
+	out := make([]map[string]any, 0, len(v.IndexingJobs))
+
+	for _, job := range v.IndexingJobs {
+		o := map[string]any{
+			"CompletedDatasources": job.CompletedDatasources,
+			"CreatedAt":            job.CreatedAt,
+			"DataSourceUuids":      job.DataSourceUuids,
+			"FinishedAt":           job.FinishedAt,
+			"KnowledgeBaseUuid":    job.KnowledgeBaseUuid,
+			"Phase":                job.Phase,
+			"StartedAt":            job.StartedAt,
+			"Status":               job.Status,
+			"Tokens":               job.Tokens,
+			"TotalDatasources":     job.TotalDatasources,
+			"TotalItemsFailed":     job.TotalItemsFailed,
+			"TotalItemsIndexed":    job.TotalItemsIndexed,
+			"TotalItemsSkipped":    job.TotalItemsSkipped,
+			"UpdatedAt":            job.UpdatedAt,
+			"UUID":                 job.Uuid,
+		}
+		out = append(out, o)
+	}
+
+	return out
+}
+
 type FunctionRoute struct {
 	Agent do.Agent
 }

commands/genai_knowledge_base.go

@@ -161,6 +161,73 @@ func KnowledgeBaseCmd() *Command {
 	cmdDataSourceDelete.Example = "The following example deletes data source having uuid like " + `00000000-0000-0000-0000-000000000000` + " from a Knowledge Base having uuid " + "`" + `f81d4fae-7dec-11d0-a765-00a0c91e6bf6` + " \nUsing the following command `" +
 		" : `doctl genai knowledge-base delete-datasource f81d4fae-7dec-11d0-a765-00a0c91e6bf6 00000000-0000-0000-0000-000000000000`"
 
+	cmdIndexingJobsList := "List all indexing jobs for knowledge bases. Each indexing job contains the following information:\n" +
+		"		- The indexing job UUID\n" +
+		"		- The knowledge base UUID\n" +
+		"		- The current phase of the job\n" +
+		"		- The job status\n" +
+		"		- The number of completed datasources\n" +
+		"		- The total number of datasources\n" +
+		"		- The number of tokens processed\n" +
+		"		- The number of items indexed\n" +
+		"		- The number of items failed\n" +
+		"		- The number of items skipped\n" +
+		"		- The creation timestamp\n" +
+		"		- The start timestamp\n" +
+		"		- The finish timestamp\n" +
+		"		- The update timestamp\n" +
+		"		- The data source UUIDs being processed"
+	cmdIndexingJobList := CmdBuilder(
+		cmd,
+		RunKnowledgeBaseListIndexingJobs,
+		"list-indexing-jobs",
+		"List all indexing jobs for knowledge bases",
+		cmdIndexingJobsList,
+		Writer, aliasOpt("ls-jobs"),
+		displayerType(&displayers.IndexingJob{}),
+	)
+	cmdIndexingJobList.Example = "The following command lists all indexing jobs for knowledge bases: " +
+		"`doctl genai knowledge-base list-indexing-jobs`"
+
+	cmdGetIndexingJobDetails := "Retrieve the status of a specific indexing job by its UUID. This includes phase, status, progress information, and timestamps."
+	cmdGetIndexingJob := CmdBuilder(
+		cmd,
+		RunKnowledgeBaseGetIndexingJob,
+		"get-indexing-job <indexing-job-uuid>",
+		"Retrieve status of indexing job for a knowledge base",
+		cmdGetIndexingJobDetails,
+		Writer, aliasOpt("get-job"),
+		displayerType(&displayers.IndexingJob{}),
+	)
+	cmdGetIndexingJob.Example = "The following command retrieves the status of an indexing job with UUID `12345678-1234-1234-1234-123456789012`: " +
+		"`doctl genai knowledge-base get-indexing-job 12345678-1234-1234-1234-123456789012`"
+
+	cmdCancelIndexingJobDetails := "Cancel a running indexing job by its UUID. This will stop the indexing process and update the job status."
+	cmdCancelIndexingJob := CmdBuilder(
+		cmd,
+		RunKnowledgeBaseCancelIndexingJob,
+		"cancel-indexing-job <indexing-job-uuid>",
+		"Cancel indexing job for a knowledge base",
+		cmdCancelIndexingJobDetails,
+		Writer, aliasOpt("cancel-job"),
+		displayerType(&displayers.IndexingJob{}),
+	)
+	cmdCancelIndexingJob.Example = "The following command cancels an indexing job with UUID `12345678-1234-1234-1234-123456789012`: " +
+		"`doctl genai knowledge-base cancel-indexing-job 12345678-1234-1234-1234-123456789012`"
+
+	cmdListIndexingJobDataSourcesDetails := "List all data sources for a specific indexing job by its UUID. This shows the status and progress of each data source being processed."
+	cmdListIndexingJobDataSources := CmdBuilder(
+		cmd,
+		RunKnowledgeBaseListIndexingJobDataSources,
+		"list-indexing-job-data-sources <indexing-job-uuid>",
+		"List data sources for indexing job for a knowledge base",
+		cmdListIndexingJobDataSourcesDetails,
+		Writer, aliasOpt("ls-job-ds"),
+		displayerType(&displayers.IndexingJobDataSource{}),
+	)
+	cmdListIndexingJobDataSources.Example = "The following command lists all data sources for an indexing job with UUID `12345678-1234-1234-1234-123456789012`: " +
+		"`doctl genai knowledge-base list-indexing-job-data-sources 12345678-1234-1234-1234-123456789012`"
+
 	cmdAttachKnowledgeBaseDetails := "Attach a knowledge base to an agent using knowledge base uuid and agent uuid. It returns the information of corresponding agent."
 	cmdAttachKnowledgeBase := CmdBuilder(
 		cmd,
@@ -462,3 +529,48 @@ func RunDetachKnowledgeBase(c *CmdConfig) error {
 		return fmt.Errorf("operation aborted")
 	}
 }
+
+// RunKnowledgeBaseListIndexingJobs lists all indexing jobs for knowledge bases.
+func RunKnowledgeBaseListIndexingJobs(c *CmdConfig) error {
+	indexingJobs, err := c.GenAI().ListIndexingJobs()
+	if err != nil {
+		return err
+	}
+	return c.Display(&displayers.IndexingJob{IndexingJobs: indexingJobs})
+}
+
+// RunKnowledgeBaseGetIndexingJob retrieves the status of a specific indexing job.
+func RunKnowledgeBaseGetIndexingJob(c *CmdConfig) error {
+	if len(c.Args) < 1 {
+		return doctl.NewMissingArgsErr(c.NS)
+	}
+	indexingJob, err := c.GenAI().GetIndexingJob(c.Args[0])
+	if err != nil {
+		return err
+	}
+	return c.Display(&displayers.IndexingJob{IndexingJobs: do.IndexingJobs{*indexingJob}})
+}
+
+// RunKnowledgeBaseCancelIndexingJob cancels a specific indexing job.
+func RunKnowledgeBaseCancelIndexingJob(c *CmdConfig) error {
+	if len(c.Args) < 1 {
+		return doctl.NewMissingArgsErr(c.NS)
+	}
+	indexingJob, err := c.GenAI().CancelIndexingJob(c.Args[0])
+	if err != nil {
+		return err
+	}
+	return c.Display(&displayers.IndexingJob{IndexingJobs: do.IndexingJobs{*indexingJob}})
+}
+
+// RunKnowledgeBaseListIndexingJobDataSources lists all data sources for a specific indexing job.
+func RunKnowledgeBaseListIndexingJobDataSources(c *CmdConfig) error {
+	if len(c.Args) < 1 {
+		return doctl.NewMissingArgsErr(c.NS)
+	}
+	dataSources, err := c.GenAI().ListIndexingJobDataSources(c.Args[0])
+	if err != nil {
+		return err
+	}
+	return c.Display(&displayers.IndexingJobDataSource{IndexingJobDataSources: dataSources})
+}

commands/genai_knowledge_base_test.go

@@ -2,6 +2,7 @@ package commands
 
 import (
 	"testing"
+	"time"
 
 	"github.com/digitalocean/doctl"
 	"github.com/digitalocean/doctl/do"
@@ -28,12 +29,48 @@ var (
 			Uuid: "data-source-id",
 		},
 	}
+
+	testIndexingJob = do.IndexingJob{
+		LastIndexingJob: &godo.LastIndexingJob{
+			CompletedDatasources: 1,
+			CreatedAt:            &godo.Timestamp{Time: time.Now()},
+			DataSourceUuids:      []string{"data-source-uuid-1", "data-source-uuid-2"},
+			FinishedAt:           &godo.Timestamp{Time: time.Now()},
+			KnowledgeBaseUuid:    "kb-uuid-123",
+			Phase:                "BATCH_JOB_PHASE_SUCCEEDED",
+			StartedAt:            &godo.Timestamp{Time: time.Now()},
+			Status:               "INDEX_JOB_STATUS_COMPLETED",
+			Tokens:               1000,
+			TotalDatasources:     2,
+			TotalItemsFailed:     "0",
+			TotalItemsIndexed:    "100",
+			TotalItemsSkipped:    "5",
+			UpdatedAt:            &godo.Timestamp{Time: time.Now()},
+			Uuid:                 "indexing-job-uuid-123",
+		},
+	}
+
+	testIndexingJobDataSource = do.IndexingJobDataSource{
+		IndexedDataSource: &godo.IndexedDataSource{
+			CompletedAt:       &godo.Timestamp{Time: time.Now()},
+			DataSourceUuid:    "data-source-uuid-1",
+			StartedAt:         &godo.Timestamp{Time: time.Now()},
+			Status:            "DATA_SOURCE_STATUS_COMPLETED",
+			IndexedItemCount:  "100",
+			FailedItemCount:   "0",
+			SkippedItemCount:  "5",
+			IndexedFileCount:  "50",
+			TotalFileCount:    "50",
+			TotalBytes:        "1024000",
+			TotalBytesIndexed: "1024000",
+		},
+	}
 )
 
 func TestKnowledgeBasesCommand(t *testing.T) {
 	cmd := KnowledgeBaseCmd()
 	assert.NotNil(t, cmd)
-	assertCommandNames(t, cmd, "add-datasource", "attach", "create", "delete", "delete-datasource", "detach", "get", "list", "list-datasources", "update")
+	assertCommandNames(t, cmd, "add-datasource", "attach", "cancel-indexing-job", "create", "delete", "delete-datasource", "detach", "get", "get-indexing-job", "list", "list-datasources", "list-indexing-job-data-sources", "list-indexing-jobs", "update")
 }
 
 func TestKnowledgeBaseGet(t *testing.T) {
@@ -214,3 +251,41 @@ func TestKnowledgeBaseDetach(t *testing.T) {
 		assert.NoError(t, err)
 	})
 }
+
+func TestKnowledgeBaseListIndexingJobs(t *testing.T) {
+	withTestClient(t, func(config *CmdConfig, tm *tcMocks) {
+		tm.genAI.EXPECT().ListIndexingJobs().Return(do.IndexingJobs{testIndexingJob}, nil)
+		err := RunKnowledgeBaseListIndexingJobs(config)
+		assert.NoError(t, err)
+	})
+}
+
+func TestKnowledgeBaseGetIndexingJob(t *testing.T) {
+	withTestClient(t, func(config *CmdConfig, tm *tcMocks) {
+		indexing_job_id := "indexing-job-uuid-123"
+		config.Args = append(config.Args, indexing_job_id)
+		tm.genAI.EXPECT().GetIndexingJob(indexing_job_id).Return(&testIndexingJob, nil)
+		err := RunKnowledgeBaseGetIndexingJob(config)
+		assert.NoError(t, err)
+	})
+}
+
+func TestKnowledgeBaseCancelIndexingJob(t *testing.T) {
+	withTestClient(t, func(config *CmdConfig, tm *tcMocks) {
+		indexing_job_id := "indexing-job-uuid-123"
+		config.Args = append(config.Args, indexing_job_id)
+		tm.genAI.EXPECT().CancelIndexingJob(indexing_job_id).Return(&testIndexingJob, nil)
+		err := RunKnowledgeBaseCancelIndexingJob(config)
+		assert.NoError(t, err)
+	})
+}
+
+func TestKnowledgeBaseListIndexingJobDataSources(t *testing.T) {
+	withTestClient(t, func(config *CmdConfig, tm *tcMocks) {
+		indexing_job_id := "indexing-job-uuid-123"
+		config.Args = append(config.Args, indexing_job_id)
+		tm.genAI.EXPECT().ListIndexingJobDataSources(indexing_job_id).Return(do.IndexingJobDataSources{testIndexingJobDataSource}, nil)
+		err := RunKnowledgeBaseListIndexingJobDataSources(config)
+		assert.NoError(t, err)
+	})
+}