From d2a2c2c361f538d98d3240e817c46d9fbec39cdb Mon Sep 17 00:00:00 2001
From: Yogesh Chauhan <yogesh.chauhan@harness.io>
Date: Mon, 10 Feb 2025 11:42:43 +0000
Subject: [PATCH] fix: [ML-648]: Add error to generate output structs (#3397)

* fix test
* pointer to message
* Add error to generate output structs
---
 app/api/controller/aiagent/types/common.go |  2 +-
 app/api/handler/aiagent/sanitize_test.go   | 12 ++++++------
 types/aigenerate/stage.go                  |  3 ++-
 types/aigenerate/step.go                   |  3 ++-
 4 files changed, 11 insertions(+), 9 deletions(-)

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 {