Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 159 additions & 0 deletions commands/displayers/genai.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
112 changes: 112 additions & 0 deletions commands/genai_knowledge_base.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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})
}
77 changes: 76 additions & 1 deletion commands/genai_knowledge_base_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package commands

import (
"testing"
"time"

"github.com/digitalocean/doctl"
"github.com/digitalocean/doctl/do"
Expand All @@ -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) {
Expand Down Expand Up @@ -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)
})
}
Loading
Loading