diff --git a/app/api/controller/aiagent/types/common.go b/app/api/controller/aiagent/types/common.go
index d2b31881b..4d8bb7f22 100644
--- a/app/api/controller/aiagent/types/common.go
+++ b/app/api/controller/aiagent/types/common.go
@@ -27,7 +27,7 @@ type Suggestion struct {
 
 type Conversation struct {
 	Role    enum.Role `json:"role"`
-	Message Message   `json:"message"`
+	Message *Message  `json:"message"`
 }
 
 type Message struct {
diff --git a/app/api/handler/aiagent/sanitize_test.go b/app/api/handler/aiagent/sanitize_test.go
index 5d3d71d18..fcb09fa21 100644
--- a/app/api/handler/aiagent/sanitize_test.go
+++ b/app/api/handler/aiagent/sanitize_test.go
@@ -47,7 +47,7 @@ func TestSanitizeConversation(t *testing.T) {
 				conversations: []types.Conversation{
 					{
 						Role: enum.RoleUser,
-						Message: types.Message{
+						Message: &types.Message{
 							Type: enum.TypeText,
 							Data: "data",
 						},
@@ -62,7 +62,7 @@ func TestSanitizeConversation(t *testing.T) {
 				conversations: []types.Conversation{
 					{
 						Role: "",
-						Message: types.Message{
+						Message: &types.Message{
 							Type: enum.TypeText,
 							Data: "data",
 						},
@@ -77,7 +77,7 @@ func TestSanitizeConversation(t *testing.T) {
 				conversations: []types.Conversation{
 					{
 						Role: "invalidRole",
-						Message: types.Message{
+						Message: &types.Message{
 							Type: enum.TypeText,
 							Data: "data",
 						},
@@ -93,7 +93,7 @@ func TestSanitizeConversation(t *testing.T) {
 				conversations: []types.Conversation{
 					{
 						Role: enum.RoleUser,
-						Message: types.Message{
+						Message: &types.Message{
 							Type: "",
 							Data: "data",
 						},
@@ -108,7 +108,7 @@ func TestSanitizeConversation(t *testing.T) {
 				conversations: []types.Conversation{
 					{
 						Role: enum.RoleUser,
-						Message: types.Message{
+						Message: &types.Message{
 							Type: "invalidMessageType",
 							Data: "data",
 						},
@@ -124,7 +124,7 @@ func TestSanitizeConversation(t *testing.T) {
 				conversations: []types.Conversation{
 					{
 						Role: enum.RoleUser,
-						Message: types.Message{
+						Message: &types.Message{
 							Type: enum.TypeText,
 							Data: "",
 						},
diff --git a/types/aigenerate/stage.go b/types/aigenerate/stage.go
index a805bcdda..68b06a46b 100644
--- a/types/aigenerate/stage.go
+++ b/types/aigenerate/stage.go
@@ -25,7 +25,8 @@ type PipelineStageGenerateRequest struct {
 }
 
 type PipelineStageGenerateResponse struct {
-	YAML string
+	YAML  string
+	Error string
 }
 
 type PipelineStageUpdateRequest struct {
diff --git a/types/aigenerate/step.go b/types/aigenerate/step.go
index aa582b27e..05054e01d 100644
--- a/types/aigenerate/step.go
+++ b/types/aigenerate/step.go
@@ -25,7 +25,8 @@ type PipelineStepGenerateRequest struct {
 }
 
 type PipelineStepGenerateResponse struct {
-	Yaml string
+	Yaml  string
+	Error string
 }
 
 type PipelineStepUpdateRequest struct {