diff --git a/app/api/controller/aiagent/types/chat.go b/app/api/controller/aiagent/types/chat.go
index 4d9699aa6..efd27d448 100644
--- a/app/api/controller/aiagent/types/chat.go
+++ b/app/api/controller/aiagent/types/chat.go
@@ -22,28 +22,28 @@ import (
 
 // Chat represents a chat conversation with a unique ID, prompt, metadata, and conversation history.
 type Chat struct {
-	ConversationID uuid.UUID         `json:"conversation_uuid,omitempty"`
-	Prompt         string            `json:"prompt"`
-	Metadata       map[string]string `json:"metadata"`
+	ConversationID uuid.UUID         `json:"conversation_id,omitempty"`
+	Prompt         string            `json:"prompt" required:"true"`
+	Metadata       map[string]string `json:"metadata" required:"true"`
 	Conversation   []Conversation    `json:"conversation"`
 }
 
 // ChatOutput represents the output of a chat request, including the conversation ID, explanation, and response.
 type ChatOutput struct {
-	ConversationID uuid.UUID      `json:"conversation_uuid"`
+	ConversationID uuid.UUID      `json:"conversation_id" required:"true"`
 	Explanation    string         `json:"explanation,omitempty"`
-	Response       []ChatResponse `json:"response"`
+	Response       []ChatResponse `json:"response" required:"true"`
 }
 
 // ChatResponse represents a response in a chat conversation, including the role and message.
 type ChatResponse struct {
-	Role    enum.Role       `json:"role"`
-	Message ResponseMessage `json:"message"`
+	Role    enum.Role       `json:"role" required:"true"`
+	Message ResponseMessage `json:"message" required:"true"`
 }
 
 // ResponseMessage represents a message in a chat response, including the type, data, and actions.
 type ResponseMessage struct {
-	Type    enum.MessageType `json:"type"`
-	Data    string           `json:"data"`
+	Type    enum.MessageType `json:"type" required:"true"`
+	Data    string           `json:"data" required:"true"`
 	Actions []enum.Action    `json:"actions,omitempty"`
 }
diff --git a/genai/genai.go b/genai/genai.go
index c64049769..1111a33b5 100644
--- a/genai/genai.go
+++ b/genai/genai.go
@@ -15,6 +15,7 @@
 package genai
 
 import (
+	"github.com/harness/gitness/app/api/controller/aiagent/types/enum"
 	capabilities2 "github.com/harness/gitness/types/capabilities"
 )
 
@@ -55,8 +56,8 @@ func (PipelineContext) GetType() capabilities2.AIContextPayloadType {
 const AIContextPayloadTypeStepContext capabilities2.AIContextPayloadType = "other"
 
 type StepContext struct {
-	Yaml      string `json:"step_yaml"`
-	StageType string `json:"stage_type"`
+	Yaml      string         `json:"step_yaml"`
+	StageType enum.StageType `json:"stage_type"`
 }
 
 func (StepContext) GetName() string {