Add extra prompt in org context

weni-ai · Jan 16, 2025 · 3cc101b · 3cc101b
1 parent 09baf41
commit 3cc101b
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 7 deletions.
diff --git a/assets/org_context.go b/assets/org_context.go
@@ -12,6 +12,7 @@ type OrgContext interface {
 	ProjectUUID() uuids.UUID
 	HasVtexAds() bool
 	HideUnavailable() bool
+	ExtraPrompt() string
 }
 
 type OrgContextReference struct {
@@ -20,10 +21,11 @@ type OrgContextReference struct {
 	ProjectUUID     uuids.UUID `json:"project_uuid"`
 	HasVtexAds      bool       `json:"vtex_ads"`
 	HideUnavailable bool       `json:"hide_unavailable"`
+	ExtraPrompt     string     `json:"extra_prompt"`
 }
 
-func NewOrgContextReference(orgContext string, projectUUID uuids.UUID, hasVtexAds bool, hideUnavailable bool) *OrgContextReference {
-	return &OrgContextReference{Context: orgContext, ProjectUUID: projectUUID, HasVtexAds: hasVtexAds, HideUnavailable: hideUnavailable}
+func NewOrgContextReference(orgContext string, projectUUID uuids.UUID, hasVtexAds bool, hideUnavailable bool, extraPrompt string) *OrgContextReference {
+	return &OrgContextReference{Context: orgContext, ProjectUUID: projectUUID, HasVtexAds: hasVtexAds, HideUnavailable: hideUnavailable, ExtraPrompt: extraPrompt}
 }
 
 func (r *OrgContextReference) Type() string {

diff --git a/assets/static/org_context.go b/assets/static/org_context.go
@@ -11,15 +11,17 @@ type OrgContext struct {
 	ProjectUUID_     uuids.UUID         `json:"project_uuid"`
 	HasVtexAds_      bool               `json:"vtex_ads"`
 	HideUnavailable_ bool               `json:"hide_unavailable"`
+	ExtraPrompt_     string             `json:"extra_prompt"`
 }
 
-func NewOrgContext(context string, channelUUID assets.ChannelUUID, projectUUID uuids.UUID, hasVtexAds bool, hideUnavailable bool) assets.OrgContext {
+func NewOrgContext(context string, channelUUID assets.ChannelUUID, projectUUID uuids.UUID, hasVtexAds bool, hideUnavailable bool, extraPrompt string) assets.OrgContext {
 	return &OrgContext{
 		Context_:         context,
 		ChannelUUID_:     channelUUID,
 		ProjectUUID_:     projectUUID,
 		HasVtexAds_:      hasVtexAds,
 		HideUnavailable_: hideUnavailable,
+		ExtraPrompt_:     extraPrompt,
 	}
 }
 
@@ -32,3 +34,5 @@ func (c *OrgContext) ProjectUUID() uuids.UUID { return c.ProjectUUID_ }
 func (c *OrgContext) HasVtexAds() bool { return c.HasVtexAds_ }
 
 func (c *OrgContext) HideUnavailable() bool { return c.HideUnavailable_ }
+
+func (c *OrgContext) ExtraPrompt() string { return c.ExtraPrompt_ }
diff --git a/assets/static/org_context_test.go b/assets/static/org_context_test.go
@@ -10,7 +10,7 @@ import (
 )
 
 func TestOrgContext(t *testing.T) {
-	oc := static.NewOrgContext("context 1", assets.ChannelUUID("e3c07fe2-7542-42c6-a394-18d968999f51"), uuids.UUID("80cb8631-b84a-4949-a483-a3103f00edc2"), false, false)
+	oc := static.NewOrgContext("context 1", assets.ChannelUUID("e3c07fe2-7542-42c6-a394-18d968999f51"), uuids.UUID("80cb8631-b84a-4949-a483-a3103f00edc2"), false, false, "")
 
 	assert.Equal(t, "context 1", oc.Context())
 	assert.Equal(t, assets.ChannelUUID("e3c07fe2-7542-42c6-a394-18d968999f51"), oc.ChannelUUID())

diff --git a/flows/org_context.go b/flows/org_context.go
@@ -67,7 +67,7 @@ func (c *OrgContext) Reference() *assets.OrgContextReference {
 	if c == nil {
 		return nil
 	}
-	return assets.NewOrgContextReference(c.Context(), c.ProjectUUID(), c.HasVtexAds(), c.HideUnavailable())
+	return assets.NewOrgContextReference(c.Context(), c.ProjectUUID(), c.HasVtexAds(), c.HideUnavailable(), c.ExtraPrompt())
 }
 
 type OrgContextAssets struct {

diff --git a/flows/org_context_test.go b/flows/org_context_test.go
@@ -11,7 +11,7 @@ import (
 )
 
 func TestOrgContext(t *testing.T) {
-	oc1 := static.NewOrgContext("context 1", assets.ChannelUUID("e3c07fe2-7542-42c6-a394-18d968999f51"), uuids.UUID("80cb8631-b84a-4949-a483-a3103f00edc2"), false, false)
+	oc1 := static.NewOrgContext("context 1", assets.ChannelUUID("e3c07fe2-7542-42c6-a394-18d968999f51"), uuids.UUID("80cb8631-b84a-4949-a483-a3103f00edc2"), false, false, "")
 
 	oc := flows.NewOrgContextAssets([]assets.OrgContext{oc1})
 
@@ -21,7 +21,7 @@ func TestOrgContext(t *testing.T) {
 	assert.Equal(t, assets.ChannelUUID("e3c07fe2-7542-42c6-a394-18d968999f51"), o1.ChannelUUID())
 	assert.Equal(t, uuids.UUID("80cb8631-b84a-4949-a483-a3103f00edc2"), o1.ProjectUUID())
 	assert.Equal(t, oc1, o1.Asset())
-	assert.Equal(t, assets.NewOrgContextReference("context 1", uuids.UUID("80cb8631-b84a-4949-a483-a3103f00edc2"), false, false), o1.Reference())
+	assert.Equal(t, assets.NewOrgContextReference("context 1", uuids.UUID("80cb8631-b84a-4949-a483-a3103f00edc2"), false, false, ""), o1.Reference())
 
 	// nil object returns nil reference
 	assert.Nil(t, (*flows.OrgContext)(nil).Reference())