Skip to content

Commit 74b9c0d

Browse files
SSharma-10SSharma-10
authored
Added KnowledgeBase Operation (#1399)
* Added KnowledgeBase Operation * Removed unused functions * format code * Added example for datasource
1 parent 40aba94 commit 74b9c0d

16 files changed

Lines changed: 1828 additions & 134 deletions

‎digitalocean/genai/datasource_knowledge_base.go‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
package genai
2+
3+
import (
4+
"context"
5+
6+
"github.com/digitalocean/terraform-provider-digitalocean/digitalocean/config"
7+
"github.com/digitalocean/terraform-provider-digitalocean/digitalocean/util"
8+
"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
9+
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
10+
)
11+
12+
func DataSourceDigitalOceanKnowledgeBase() *schema.Resource {
13+
14+
return &schema.Resource{
15+
ReadContext: dataSourceDigitalOceanKnowledgeBasesRead,
16+
Schema: KnowledgeBaseSchemaRead(),
17+
}
18+
}
19+
20+
func DataSourceDigitalOceanKnowledgeBaseDatasources() *schema.Resource {
21+
return &schema.Resource{
22+
ReadContext: dataSourceDigitalOceanKnowledgeBaseDatasourcesRead,
23+
Schema: map[string]*schema.Schema{
24+
"knowledge_base_uuid": {
25+
Type: schema.TypeString,
26+
Required: true,
27+
Description: "UUID of the Knowledge Base",
28+
},
29+
"datasources": {
30+
Type: schema.TypeList,
31+
Computed: true,
32+
Description: "List of data sources for the Knowledge Base",
33+
Elem: knowledgeBaseDatasourcesSchema(),
34+
},
35+
},
36+
}
37+
}
38+
39+
func dataSourceDigitalOceanKnowledgeBasesRead(ctx context.Context, d *schema.ResourceData, meta interface{}) diag.Diagnostics {
40+
client := meta.(*config.CombinedConfig).GodoClient()
41+
kbIDRaw, ok := d.GetOk("uuid")
42+
if !ok || kbIDRaw == nil {
43+
return diag.Errorf("uuid must be provided")
44+
}
45+
kbID := kbIDRaw.(string)
46+
47+
kb, _, _, err := client.GenAI.GetKnowledgeBase(ctx, kbID)
48+
if err != nil {
49+
return diag.FromErr(err)
50+
}
51+
52+
flattened, err := FlattenDigitalOceanKnowledgeBase(kb)
53+
if err != nil {
54+
return diag.FromErr(err)
55+
}
56+
57+
if err := util.SetResourceDataFromMap(d, flattened); err != nil {
58+
return diag.FromErr(err)
59+
}
60+
61+
d.SetId(kb.Uuid)
62+
return nil
63+
}
64+
65+
func dataSourceDigitalOceanKnowledgeBaseDatasourcesRead(ctx context.Context, d *schema.ResourceData, meta interface{}) diag.Diagnostics {
66+
client := meta.(*config.CombinedConfig).GodoClient()
67+
68+
kbUUIDRaw, ok := d.GetOk("knowledge_base_uuid")
69+
if !ok || kbUUIDRaw == nil {
70+
return diag.Errorf("knowledge_base_uuid must be provided")
71+
}
72+
kbUUID := kbUUIDRaw.(string)
73+
74+
// Call the API to list data sources for the KB
75+
datasources, _, err := client.GenAI.ListKnowledgeBaseDataSources(ctx, kbUUID, nil)
76+
if err != nil {
77+
return diag.FromErr(err)
78+
}
79+
80+
// Flatten and set datasources
81+
flattened := flattenKnowledgeBaseDataSources(datasources)
82+
if err := d.Set("datasources", flattened); err != nil {
83+
return diag.FromErr(err)
84+
}
85+
86+
d.SetId(kbUUID)
87+
return nil
88+
}

‎digitalocean/genai/datasource_knowledge_base_test.go‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
package genai_test
2+
3+
import (
4+
"fmt"
5+
"testing"
6+
7+
"github.com/digitalocean/terraform-provider-digitalocean/digitalocean/acceptance"
8+
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
9+
)
10+
11+
func TestAccDataSourceDigitalOceanKnowledgeBase_BasicByID(t *testing.T) {
12+
kbName := acceptance.RandomTestName() + "-kb"
13+
resourceConfig := fmt.Sprintf(`
14+
resource "digitalocean_genai_knowledge_base" "test" {
15+
name = "%s"
16+
project_id = "%s"
17+
region = "%s"
18+
embedding_model_uuid = "%s"
19+
tags = ["terraform-test", "datasource-test"]
20+
is_public = false
21+
22+
datasources {
23+
web_crawler_data_source {
24+
base_url = "https://docs.digitalocean.com/products/kubernetes/"
25+
crawling_option = "SCOPED"
26+
embed_media = true
27+
}
28+
}
29+
}
30+
31+
data "digitalocean_genai_knowledge_base" "byid" {
32+
uuid = digitalocean_genai_knowledge_base.test.id
33+
}
34+
`, kbName, defaultProjectID, defaultRegion, defaultEmbeddingModelUUID)
35+
36+
resource.ParallelTest(t, resource.TestCase{
37+
PreCheck: func() { acceptance.TestAccPreCheck(t) },
38+
ProviderFactories: acceptance.TestAccProviderFactories,
39+
Steps: []resource.TestStep{
40+
{
41+
Config: resourceConfig,
42+
Check: resource.ComposeTestCheckFunc(
43+
resource.TestCheckResourceAttrPair(
44+
"data.digitalocean_genai_knowledge_base.byid", "id",
45+
"digitalocean_genai_knowledge_base.test", "id",
46+
),
47+
resource.TestCheckResourceAttr(
48+
"data.digitalocean_genai_knowledge_base.byid", "name", kbName),
49+
resource.TestCheckResourceAttr(
50+
"data.digitalocean_genai_knowledge_base.byid", "project_id", defaultProjectID),
51+
resource.TestCheckResourceAttr(
52+
"data.digitalocean_genai_knowledge_base.byid", "region", defaultRegion),
53+
resource.TestCheckResourceAttr(
54+
"data.digitalocean_genai_knowledge_base.byid", "embedding_model_uuid", defaultEmbeddingModelUUID),
55+
resource.TestCheckResourceAttr(
56+
"data.digitalocean_genai_knowledge_base.byid", "is_public", "false"),
57+
resource.TestCheckResourceAttr(
58+
"data.digitalocean_genai_knowledge_base.byid", "tags.#", "2"),
59+
resource.TestCheckResourceAttr(
60+
"data.digitalocean_genai_knowledge_base.byid", "datasources.#", "1"),
61+
resource.TestCheckResourceAttr(
62+
"data.digitalocean_genai_knowledge_base.byid", "datasources.0.web_crawler_data_source.0.base_url", "https://docs.digitalocean.com/products/kubernetes/"),
63+
resource.TestCheckResourceAttrSet(
64+
"data.digitalocean_genai_knowledge_base.byid", "created_at"),
65+
),
66+
},
67+
},
68+
})
69+
}

‎digitalocean/genai/datasource_knowledge_bases.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package genai
2+
3+
import (
4+
"github.com/digitalocean/terraform-provider-digitalocean/internal/datalist"
5+
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
6+
)
7+
8+
func DataSourceDigitalOceanKnowledgeBases() *schema.Resource {
9+
dataListConfig := &datalist.ResourceConfig{
10+
RecordSchema: KnowledgeBaseSchemaRead(),
11+
ResultAttributeName: "knowledge_bases",
12+
FlattenRecord: flattenDigitalOceanKnowledgeBase,
13+
GetRecords: getDigitalOceanKnowledgeBases,
14+
}
15+
16+
return datalist.NewResource(dataListConfig)
17+
}

‎digitalocean/genai/datasource_knowledge_bases_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package genai_test
2+
3+
import (
4+
"fmt"
5+
"testing"
6+
7+
"github.com/digitalocean/terraform-provider-digitalocean/digitalocean/acceptance"
8+
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
9+
)
10+
11+
const (
12+
defaultProjectID = "84e1e297-ee40-41ac-95ff-1067cf2206e9"
13+
defaultEmbeddingModelUUID = "22653204-79ed-11ef-bf8f-4e013e2ddde4"
14+
defaultRegion = "tor1"
15+
)
16+
17+
func TestAccDataSourceDigitalOceanKnowledgeBases_Basic(t *testing.T) {
18+
kbName1 := acceptance.RandomTestName() + "-kb1"
19+
kbName2 := acceptance.RandomTestName() + "-kb2"
20+
21+
resourceConfig := fmt.Sprintf(`
22+
resource "digitalocean_genai_knowledge_base" "test1" {
23+
name = "%s"
24+
project_id = "%s"
25+
region = "%s"
26+
embedding_model_uuid = "%s"
27+
tags = ["terraform-test", "datasource-test"]
28+
is_public = false
29+
30+
datasources {
31+
web_crawler_data_source {
32+
base_url = "https://docs.digitalocean.com/products/kubernetes/"
33+
crawling_option = "SCOPED"
34+
embed_media = true
35+
}
36+
}
37+
}
38+
39+
resource "digitalocean_genai_knowledge_base" "test2" {
40+
name = "%s"
41+
project_id = "%s"
42+
region = "%s"
43+
embedding_model_uuid = "%s"
44+
tags = ["terraform-test", "datasource-test"]
45+
is_public = true
46+
47+
datasources {
48+
web_crawler_data_source {
49+
base_url = "https://docs.digitalocean.com/products/app-platform/"
50+
crawling_option = "SCOPED"
51+
embed_media = false
52+
}
53+
}
54+
}
55+
56+
data "digitalocean_genai_knowledge_bases" "all" {}
57+
`, kbName1, defaultProjectID, defaultRegion, defaultEmbeddingModelUUID,
58+
kbName2, defaultProjectID, defaultRegion, defaultEmbeddingModelUUID)
59+
60+
resource.ParallelTest(t, resource.TestCase{
61+
PreCheck: func() { acceptance.TestAccPreCheck(t) },
62+
ProviderFactories: acceptance.TestAccProviderFactories,
63+
Steps: []resource.TestStep{
64+
{
65+
Config: resourceConfig,
66+
Check: resource.ComposeTestCheckFunc(
67+
resource.TestCheckResourceAttrSet("data.digitalocean_genai_knowledge_bases.all", "knowledge_bases.#"),
68+
),
69+
},
70+
},
71+
})
72+
}

0 commit comments

Comments
 (0)